Kaynağa Gözat

[Firebase AI] Add `URLContext` tool (#15221)

Co-authored-by: Paul Beusterien <paulbeusterien@google.com>
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Andrew Heard 6 ay önce
ebeveyn
işleme
218716f556

+ 2 - 0
FirebaseAI/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [feature] Added support for the URL context tool, which allows the model to access content
+  from provided public web URLs to inform and enhance its responses. (#15221)
 - [changed] Using Firebase AI Logic with the Gemini Developer API is now Generally Available (GA).
 - [changed] Using Firebase AI Logic with the Imagen generation APIs is now Generally Available (GA).
 

+ 1 - 0
FirebaseAI/Sources/AILog.swift

@@ -66,6 +66,7 @@ enum AILog {
     case codeExecutionResultUnrecognizedOutcome = 3015
     case executableCodeUnrecognizedLanguage = 3016
     case fallbackValueUsed = 3017
+    case urlMetadataUnrecognizedURLRetrievalStatus = 3018
 
     // SDK State Errors
     case generateContentResponseNoCandidates = 4000

+ 30 - 2
FirebaseAI/Sources/GenerateContentResponse.swift

@@ -26,6 +26,9 @@ public struct GenerateContentResponse: Sendable {
     /// The total number of tokens across the generated response candidates.
     public let candidatesTokenCount: Int
 
+    /// The number of tokens used by tools.
+    public let toolUsePromptTokenCount: Int
+
     /// The number of tokens used by the model's internal "thinking" process.
     ///
     /// For models that support thinking (like Gemini 2.5 Pro and Flash), this represents the actual
@@ -39,11 +42,15 @@ public struct GenerateContentResponse: Sendable {
     /// The total number of tokens in both the request and response.
     public let totalTokenCount: Int
 
-    /// The breakdown, by modality, of how many tokens are consumed by the prompt
+    /// The breakdown, by modality, of how many tokens are consumed by the prompt.
     public let promptTokensDetails: [ModalityTokenCount]
 
     /// The breakdown, by modality, of how many tokens are consumed by the candidates
     public let candidatesTokensDetails: [ModalityTokenCount]
+
+    /// The breakdown, by modality, of how many tokens were consumed by the tools used to process
+    /// the request.
+    public let toolUsePromptTokensDetails: [ModalityTokenCount]
   }
 
   /// A list of candidate response content, ordered from best to worst.
@@ -154,14 +161,19 @@ public struct Candidate: Sendable {
 
   public let groundingMetadata: GroundingMetadata?
 
+  /// Metadata related to the ``URLContext`` tool.
+  public let urlContextMetadata: URLContextMetadata?
+
   /// Initializer for SwiftUI previews or tests.
   public init(content: ModelContent, safetyRatings: [SafetyRating], finishReason: FinishReason?,
-              citationMetadata: CitationMetadata?, groundingMetadata: GroundingMetadata? = nil) {
+              citationMetadata: CitationMetadata?, groundingMetadata: GroundingMetadata? = nil,
+              urlContextMetadata: URLContextMetadata? = nil) {
     self.content = content
     self.safetyRatings = safetyRatings
     self.finishReason = finishReason
     self.citationMetadata = citationMetadata
     self.groundingMetadata = groundingMetadata
+    self.urlContextMetadata = urlContextMetadata
   }
 
   // Returns `true` if the candidate contains no information that a developer could use.
@@ -469,10 +481,12 @@ extension GenerateContentResponse.UsageMetadata: Decodable {
   enum CodingKeys: CodingKey {
     case promptTokenCount
     case candidatesTokenCount
+    case toolUsePromptTokenCount
     case thoughtsTokenCount
     case totalTokenCount
     case promptTokensDetails
     case candidatesTokensDetails
+    case toolUsePromptTokensDetails
   }
 
   public init(from decoder: any Decoder) throws {
@@ -480,6 +494,8 @@ extension GenerateContentResponse.UsageMetadata: Decodable {
     promptTokenCount = try container.decodeIfPresent(Int.self, forKey: .promptTokenCount) ?? 0
     candidatesTokenCount =
       try container.decodeIfPresent(Int.self, forKey: .candidatesTokenCount) ?? 0
+    toolUsePromptTokenCount =
+      try container.decodeIfPresent(Int.self, forKey: .toolUsePromptTokenCount) ?? 0
     thoughtsTokenCount = try container.decodeIfPresent(Int.self, forKey: .thoughtsTokenCount) ?? 0
     totalTokenCount = try container.decodeIfPresent(Int.self, forKey: .totalTokenCount) ?? 0
     promptTokensDetails =
@@ -488,6 +504,9 @@ extension GenerateContentResponse.UsageMetadata: Decodable {
       [ModalityTokenCount].self,
       forKey: .candidatesTokensDetails
     ) ?? []
+    toolUsePromptTokensDetails = try container.decodeIfPresent(
+      [ModalityTokenCount].self, forKey: .toolUsePromptTokensDetails
+    ) ?? []
   }
 }
 
@@ -499,6 +518,7 @@ extension Candidate: Decodable {
     case finishReason
     case citationMetadata
     case groundingMetadata
+    case urlContextMetadata
   }
 
   /// Initializes a response from a decoder. Used for decoding server responses; not for public
@@ -540,6 +560,14 @@ extension Candidate: Decodable {
       GroundingMetadata.self,
       forKey: .groundingMetadata
     )
+
+    if let urlContextMetadata =
+      try container.decodeIfPresent(URLContextMetadata.self, forKey: .urlContextMetadata),
+      !urlContextMetadata.urlMetadata.isEmpty {
+      self.urlContextMetadata = urlContextMetadata
+    } else {
+      urlContextMetadata = nil
+    }
   }
 }
 

+ 12 - 0
FirebaseAI/Sources/Tool.swift

@@ -76,12 +76,15 @@ public struct Tool: Sendable {
   let googleSearch: GoogleSearch?
 
   let codeExecution: CodeExecution?
+  let urlContext: URLContext?
 
   init(functionDeclarations: [FunctionDeclaration]? = nil,
        googleSearch: GoogleSearch? = nil,
+       urlContext: URLContext? = nil,
        codeExecution: CodeExecution? = nil) {
     self.functionDeclarations = functionDeclarations
     self.googleSearch = googleSearch
+    self.urlContext = urlContext
     self.codeExecution = codeExecution
   }
 
@@ -128,6 +131,15 @@ public struct Tool: Sendable {
     return self.init(googleSearch: googleSearch)
   }
 
+  /// Creates a tool that allows you to provide additional context to the models in the form of
+  /// public web URLs.
+  ///
+  /// By including URLs in your request, the Gemini model will access the content from those pages
+  /// to inform and enhance its response.
+  public static func urlContext() -> Tool {
+    return self.init(urlContext: URLContext())
+  }
+
   /// Creates a tool that allows the model to execute code.
   ///
   /// For more details, see ``CodeExecution``.

+ 18 - 0
FirebaseAI/Sources/Types/Internal/Tools/URLContext.swift

@@ -0,0 +1,18 @@
+// Copyright 2025 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.
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct URLContext: Sendable, Encodable {
+  init() {}
+}

+ 34 - 0
FirebaseAI/Sources/Types/Public/URLContextMetadata.swift

@@ -0,0 +1,34 @@
+// Copyright 2025 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.
+
+/// Metadata related to the ``Tool/urlContext()`` tool.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct URLContextMetadata: Sendable, Hashable {
+  /// List of URL metadata used to provide context to the Gemini model.
+  public let urlMetadata: [URLMetadata]
+}
+
+// MARK: - Codable Conformances
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension URLContextMetadata: Decodable {
+  enum CodingKeys: CodingKey {
+    case urlMetadata
+  }
+
+  public init(from decoder: any Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    urlMetadata = try container.decodeIfPresent([URLMetadata].self, forKey: .urlMetadata) ?? []
+  }
+}

+ 85 - 0
FirebaseAI/Sources/Types/Public/URLMetadata.swift

@@ -0,0 +1,85 @@
+// Copyright 2025 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.
+
+import Foundation
+
+/// Metadata for a single URL retrieved by the ``Tool/urlContext()`` tool.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct URLMetadata: Sendable, Hashable {
+  /// Status of the URL retrieval.
+  public struct URLRetrievalStatus: DecodableProtoEnum, Hashable {
+    enum Kind: String {
+      case unspecified = "URL_RETRIEVAL_STATUS_UNSPECIFIED"
+      case success = "URL_RETRIEVAL_STATUS_SUCCESS"
+      case error = "URL_RETRIEVAL_STATUS_ERROR"
+      case paywall = "URL_RETRIEVAL_STATUS_PAYWALL"
+      case unsafe = "URL_RETRIEVAL_STATUS_UNSAFE"
+    }
+
+    /// Internal only - Unspecified retrieval status.
+    static let unspecified = URLRetrievalStatus(kind: .unspecified)
+
+    /// The URL retrieval was successful.
+    public static let success = URLRetrievalStatus(kind: .success)
+
+    /// The URL retrieval failed.
+    public static let error = URLRetrievalStatus(kind: .error)
+
+    /// The URL retrieval failed because the content is behind a paywall.
+    public static let paywall = URLRetrievalStatus(kind: .paywall)
+
+    /// The URL retrieval failed because the content is unsafe.
+    public static let unsafe = URLRetrievalStatus(kind: .unsafe)
+
+    /// Returns the raw string representation of the `URLRetrievalStatus` value.
+    public let rawValue: String
+
+    static let unrecognizedValueMessageCode =
+      AILog.MessageCode.urlMetadataUnrecognizedURLRetrievalStatus
+  }
+
+  /// The retrieved URL.
+  public let retrievedURL: URL?
+
+  /// The status of the URL retrieval.
+  public let retrievalStatus: URLRetrievalStatus
+}
+
+// MARK: - Codable Conformances
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension URLMetadata: Decodable {
+  enum CodingKeys: String, CodingKey {
+    case retrievedURL = "retrievedUrl"
+    case retrievalStatus = "urlRetrievalStatus"
+  }
+
+  public init(from decoder: any Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+
+    if let retrievedURLString = try container.decodeIfPresent(String.self, forKey: .retrievedURL),
+       let retrievedURL = URL(string: retrievedURLString) {
+      self.retrievedURL = retrievedURL
+    } else {
+      retrievedURL = nil
+    }
+    let retrievalStatus = try container.decodeIfPresent(
+      URLMetadata.URLRetrievalStatus.self, forKey: .retrievalStatus
+    )
+
+    self.retrievalStatus = AILog.safeUnwrap(
+      retrievalStatus, fallback: URLMetadata.URLRetrievalStatus(kind: .unspecified)
+    )
+  }
+}

+ 27 - 0
FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

@@ -424,6 +424,33 @@ struct GenerateContentIntegrationTests {
     }
   }
 
+  @Test(
+    "generateContent with URL Context",
+    arguments: InstanceConfig.allConfigs
+  )
+  func generateContent_withURLContext_succeeds(_ config: InstanceConfig) async throws {
+    let model = FirebaseAI.componentInstance(config).generativeModel(
+      modelName: ModelNames.gemini2_5_Flash,
+      tools: [.urlContext()]
+    )
+    let prompt = """
+    Write a one paragraph summary of this blog post: \
+    https://developers.googleblog.com/en/introducing-gemma-3-270m/
+    """
+
+    let response = try await model.generateContent(prompt)
+
+    let candidate = try #require(response.candidates.first)
+    let urlContextMetadata = try #require(candidate.urlContextMetadata)
+    #expect(urlContextMetadata.urlMetadata.count == 1)
+    let urlMetadata = try #require(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try #require(urlMetadata.retrievedURL)
+    #expect(
+      retrievedURL == URL(string: "https://developers.googleblog.com/en/introducing-gemma-3-270m/")
+    )
+    #expect(urlMetadata.retrievalStatus == .success)
+  }
+
   @Test(arguments: InstanceConfig.allConfigs)
   func generateContent_codeExecution_succeeds(_ config: InstanceConfig) async throws {
     let model = FirebaseAI.componentInstance(config).generativeModel(

+ 72 - 0
FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift

@@ -333,6 +333,55 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     let textPart = try XCTUnwrap(parts[2] as? TextPart)
     XCTAssertFalse(textPart.isThought)
     XCTAssertTrue(textPart.text.hasPrefix("The first 5 prime numbers are 2, 3, 5, 7, and 11."))
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)
+  }
+
+  func testGenerateContent_success_urlContext() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-url-context",
+      withExtension: "json",
+      subdirectory: googleAISubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
+    let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
+    XCTAssertEqual(
+      retrievedURL,
+      URL(string: "https://berkshirehathaway.com")
+    )
+    XCTAssertEqual(urlMetadata.retrievalStatus, .success)
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 424)
+  }
+
+  func testGenerateContent_success_urlContext_mixedValidity() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-url-context-mixed-validity",
+      withExtension: "json",
+      subdirectory: googleAISubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)
+
+    let paywallURLMetadata = urlContextMetadata.urlMetadata[0]
+    XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)
+
+    let successURLMetadata = urlContextMetadata.urlMetadata[1]
+    XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
+
+    let errorURLMetadata = urlContextMetadata.urlMetadata[2]
+    XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
   }
 
   func testGenerateContent_failure_invalidAPIKey() async throws {
@@ -642,4 +691,27 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     let lastResponse = try XCTUnwrap(responses.last)
     XCTAssertEqual(lastResponse.text, "text8")
   }
+
+  func testGenerateContentStream_success_urlContext() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "streaming-success-url-context",
+      withExtension: "txt",
+      subdirectory: googleAISubdirectory
+    )
+
+    var responses = [GenerateContentResponse]()
+    let stream = try model.generateContentStream(testPrompt)
+    for try await response in stream {
+      responses.append(response)
+    }
+
+    let firstResponse = try XCTUnwrap(responses.first)
+    let candidate = try XCTUnwrap(firstResponse.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
+    let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
+    XCTAssertEqual(retrievedURL, URL(string: "https://google.com"))
+    XCTAssertEqual(urlMetadata.retrievalStatus, .success)
+  }
 }

+ 90 - 0
FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift

@@ -487,6 +487,73 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertEqual(
       textPart2.text, "The sum of the first 5 prime numbers (2, 3, 5, 7, and 11) is 28."
     )
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371)
+  }
+
+  func testGenerateContent_success_urlContext() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-url-context",
+      withExtension: "json",
+      subdirectory: vertexSubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
+    let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
+    XCTAssertEqual(
+      retrievedURL,
+      URL(string: "https://berkshirehathaway.com")
+    )
+    XCTAssertEqual(urlMetadata.retrievalStatus, .success)
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 34)
+    XCTAssertEqual(usageMetadata.thoughtsTokenCount, 36)
+  }
+
+  func testGenerateContent_success_urlContext_mixedValidity() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-url-context-mixed-validity",
+      withExtension: "json",
+      subdirectory: vertexSubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)
+
+    let paywallURLMetadata = urlContextMetadata.urlMetadata[0]
+    XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)
+
+    let successURLMetadata = urlContextMetadata.urlMetadata[1]
+    XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
+
+    let errorURLMetadata = urlContextMetadata.urlMetadata[2]
+    XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
+  }
+
+  func testGenerateContent_success_urlContext_retrievedURLPresentOnErrorStatus() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-url-context-missing-retrievedurl",
+      withExtension: "json",
+      subdirectory: vertexSubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
+    XCTAssertEqual(retrievedURL.absoluteString, "https://example.com/8")
+    XCTAssertEqual(urlMetadata.retrievalStatus, .error)
   }
 
   func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {
@@ -1718,6 +1785,29 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertEqual(responses, 1)
   }
 
+  func testGenerateContentStream_success_urlContext() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "streaming-success-url-context",
+      withExtension: "txt",
+      subdirectory: vertexSubdirectory
+    )
+
+    var responses = [GenerateContentResponse]()
+    let stream = try model.generateContentStream(testPrompt)
+    for try await response in stream {
+      responses.append(response)
+    }
+
+    let firstResponse = try XCTUnwrap(responses.first)
+    let candidate = try XCTUnwrap(firstResponse.candidates.first)
+    let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
+    XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
+    let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
+    let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
+    XCTAssertEqual(retrievedURL, URL(string: "https://google.com"))
+    XCTAssertEqual(urlMetadata.retrievalStatus, .success)
+  }
+
   // MARK: - Count Tokens
 
   func testCountTokens_succeeds() async throws {

+ 5 - 3
FirebaseAI/Tests/Unit/MockURLProtocol.swift

@@ -21,6 +21,7 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
     URLResponse,
     AsyncLineSequence<URL.AsyncBytes>?
   ))?
+
   override class func canInit(with request: URLRequest) -> Bool {
     #if os(watchOS)
       print("MockURLProtocol cannot be used on watchOS.")
@@ -33,13 +34,14 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
   override class func canonicalRequest(for request: URLRequest) -> URLRequest { return request }
 
   override func startLoading() {
-    guard let requestHandler = MockURLProtocol.requestHandler else {
-      fatalError("`requestHandler` is nil.")
-    }
     guard let client = client else {
       fatalError("`client` is nil.")
     }
 
+    guard let requestHandler = MockURLProtocol.requestHandler else {
+      fatalError("No request handler set.")
+    }
+
     Task {
       let (response, stream) = try requestHandler(self.request)
       client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)

+ 51 - 0
FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

@@ -17,6 +17,8 @@ import XCTest
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 final class GenerateContentResponseTests: XCTestCase {
+  let jsonDecoder = JSONDecoder()
+
   // MARK: - GenerateContentResponse Computed Properties
 
   func testGenerateContentResponse_inlineDataParts_success() throws {
@@ -106,4 +108,53 @@ final class GenerateContentResponseTests: XCTestCase {
       "functionCalls should be empty when there are no candidates."
     )
   }
+
+  // MARK: - Decoding Tests
+
+  func testDecodeCandidate_emptyURLMetadata_urlContextMetadataIsNil() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "STOP",
+      "urlContextMetadata": { "urlMetadata": [] }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertNil(
+      candidate.urlContextMetadata,
+      "urlContextMetadata should be nil if the `urlMetadata` array is empty in the candidate."
+    )
+    XCTAssertEqual(candidate.content.role, "model")
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    let textPart = try XCTUnwrap(part as? TextPart)
+    XCTAssertEqual(textPart.text, "Some text.")
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(candidate.finishReason, .stop)
+  }
+
+  func testDecodeCandidate_missingURLMetadata_urlContextMetadataIsNil() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "STOP"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertNil(
+      candidate.urlContextMetadata,
+      "urlContextMetadata should be nil if `urlMetadata` is not provided in the candidate."
+    )
+    XCTAssertEqual(candidate.content.role, "model")
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    let textPart = try XCTUnwrap(part as? TextPart)
+    XCTAssertEqual(textPart.text, "Some text.")
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(candidate.finishReason, .stop)
+  }
 }