[Firebase AI] Add support for direct Vertex AI integration testing

FirebaseAI/Sources/Constants.swift

@@ -21,4 +21,15 @@ enum Constants {
   /// - Important: A suffix must be appended to produce an error domain (e.g.,
   ///   "com.google.firebase.firebaseai.ExampleError").
   static let baseErrorDomain = "com.google.firebase.firebaseai"
+
+  #if DEBUG
+    /// The key for an environment variable containing a Google Cloud Access Token.
+    ///
+    /// This should only be used for SDK development and testing with the Vertex AI direct backend
+    /// that bypasses the Firebase proxy..
+    ///
+    /// The value should is typically obtained from the gcloud CLI by calling
+    /// `gcloud auth print-access-token`.
+    static let gCloudAccessTokenEnvVarKey = "FIRGCloudAuthAccessToken"
+  #endif // DEBUG
 }

FirebaseAI/Sources/FirebaseAI.swift

@@ -299,11 +299,18 @@ public final class FirebaseAI: Sendable {
 
   private func developerModelResourceName(modelName: String) -> String {
     switch apiConfig.service.endpoint {
-    case .firebaseProxyStaging, .firebaseProxyProd:
-      let projectID = firebaseInfo.projectID
-      return "projects/\(projectID)/models/\(modelName)"
-    case .googleAIBypassProxy:
-      return "models/\(modelName)"
+    case .firebaseProxyProd:
+      return "projects/\(firebaseInfo.projectID)/models/\(modelName)"
+    #if DEBUG
+      case .googleAIBypassProxy:
+        return "models/\(modelName)"
+      case .firebaseProxyStaging:
+        return "projects/\(firebaseInfo.projectID)/models/\(modelName)"
+      case .vertexAIStagingBypassProxy:
+        fatalError(
+          "The Vertex AI staging endpoint does not support the Gemini Developer API (Google AI)."
+        )
+    #endif // DEBUG
     }
   }
 

FirebaseAI/Sources/GenerativeAIService.swift

@@ -169,7 +169,16 @@ struct GenerativeAIService {
   private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
     var urlRequest = try URLRequest(url: request.getURL())
     urlRequest.httpMethod = "POST"
-    urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key")
+    #if DEBUG
+      let accessToken = ProcessInfo.processInfo.environment[Constants.gCloudAccessTokenEnvVarKey]
+    #else
+      let accessToken: String? = nil
+    #endif // DEBUG
+    if let accessToken {
+      urlRequest.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
+    } else {
+      urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key")
+    }
     urlRequest.setValue(
       "\(GenerativeAIService.languageTag) \(GenerativeAIService.firebaseVersionTag)",
       forHTTPHeaderField: "x-goog-api-client"
@@ -190,9 +199,8 @@ struct GenerativeAIService {
       }
     }
 
-    if let auth = firebaseInfo.auth, let authToken = try await auth.getToken(
-      forcingRefresh: false
-    ) {
+    if let auth = firebaseInfo.auth, let authToken = try await auth.getToken(forcingRefresh: false),
+       accessToken == nil {
       urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization")
     }
 

FirebaseAI/Sources/GenerativeModel.swift

@@ -322,11 +322,20 @@ public final class GenerativeModel: Sendable {
     // "models/model-name". This field is unaltered by the Firebase backend before forwarding the
     // request to the Generative Language backend, which expects the form "models/model-name".
     let generateContentRequestModelResourceName = switch apiConfig.service {
-    case .vertexAI, .googleAI(endpoint: .googleAIBypassProxy):
+    case .vertexAI:
       modelResourceName
-    case .googleAI(endpoint: .firebaseProxyProd),
-         .googleAI(endpoint: .firebaseProxyStaging):
+    case .googleAI(endpoint: .firebaseProxyProd):
       "models/\(modelName)"
+    #if DEBUG
+      case .googleAI(endpoint: .firebaseProxyStaging):
+        "models/\(modelName)"
+      case .googleAI(endpoint: .googleAIBypassProxy):
+        modelResourceName
+      case .googleAI(endpoint: .vertexAIStagingBypassProxy):
+        fatalError(
+          "The Vertex AI staging endpoint does not support the Gemini Developer API (Google AI)."
+        )
+    #endif // DEBUG
     }
 
     let generateContentRequest = GenerateContentRequest(

FirebaseAI/Sources/Types/Internal/APIConfig.swift

@@ -75,28 +75,43 @@ extension APIConfig.Service {
     /// This endpoint supports both Google AI and Vertex AI.
     case firebaseProxyProd = "https://firebasevertexai.googleapis.com"
 
-    /// The Firebase proxy staging endpoint; for SDK development and testing only.
-    ///
-    /// This endpoint supports both the Gemini Developer API (commonly referred to as Google AI)
-    /// and the Gemini API in Vertex AI (commonly referred to simply as Vertex AI).
-    case firebaseProxyStaging = "https://staging-firebasevertexai.sandbox.googleapis.com"
+    #if DEBUG
+      /// The Firebase proxy staging endpoint; for SDK development and testing only.
+      ///
+      /// This endpoint supports both the Gemini Developer API (commonly referred to as Google AI)
+      /// and the Gemini API in Vertex AI (commonly referred to simply as Vertex AI).
+      case firebaseProxyStaging = "https://staging-firebasevertexai.sandbox.googleapis.com"
 
-    /// The Gemini Developer API (Google AI) direct production endpoint; for SDK development and
-    /// testing only.
-    ///
-    /// This bypasses the Firebase proxy and directly connects to the Gemini Developer API
-    /// (Google AI) backend. This endpoint only supports the Gemini Developer API, not Vertex AI.
-    case googleAIBypassProxy = "https://generativelanguage.googleapis.com"
+      /// The Gemini Developer API (Google AI) direct production endpoint; for SDK development and
+      /// testing only.
+      ///
+      /// This bypasses the Firebase proxy and directly connects to the Gemini Developer API
+      /// (Google AI) backend. This endpoint only supports the Gemini Developer API, not Vertex AI.
+      case googleAIBypassProxy = "https://generativelanguage.googleapis.com"
+
+      /// The Vertex AI direct staging endpoint; for SDK development and testing only.
+      ///
+      /// This bypasses the Firebase proxy and directly connects to the Vertex AI backend. This
+      /// endpoint only supports the Gemini API in Vertex AI, not the Gemini Developer API.
+      case vertexAIStagingBypassProxy = "https://staging-aiplatform.sandbox.googleapis.com"
+    #endif // DEBUG
   }
 }
 
 extension APIConfig {
   /// Versions of the configured API service (`APIConfig.Service`).
   enum Version: String, Encodable {
-    /// The stable channel for version 1 of the API.
-    case v1
-
     /// The beta channel for version 1 of the API.
     case v1beta
+
+    #if DEBUG
+      /// The stable channel for version 1 of the API; currently for SDK development and testing
+      /// only.
+      case v1
+
+      /// The beta channel for version 1 of the direct Vertex AI API, when bypassing the Firebase
+      /// proxy; for SDK development and testing only.
+      case v1beta1
+    #endif // DEBUG
   }
 }

FirebaseAI/Tests/TestApp/FirebaseAITestApp.xcodeproj/xcshareddata/xcschemes/FirebaseAITestApp-SPM.xcscheme

@@ -69,6 +69,13 @@
             isEnabled = "YES">
          </CommandLineArgument>
       </CommandLineArguments>
+      <EnvironmentVariables>
+         <EnvironmentVariable
+            key = "FIRGCloudAuthAccessToken"
+            value = "Run `gcloud auth print-access-token` to obtain as access token for direct Vertex AI testing."
+            isEnabled = "NO">
+         </EnvironmentVariable>
+      </EnvironmentVariables>
    </LaunchAction>
    <ProfileAction
       buildConfiguration = "Release"

FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift

@@ -52,6 +52,12 @@ struct InstanceConfig: Equatable, Encodable {
       version: .v1beta
     )
   )
+  static let vertexAI_v1beta_staging_global_bypassProxy = InstanceConfig(
+    apiConfig: APIConfig(
+      service: .vertexAI(endpoint: .vertexAIStagingBypassProxy, location: "global"),
+      version: .v1beta1
+    )
+  )
   static let googleAI_v1beta = InstanceConfig(
     apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
   )
@@ -80,6 +86,7 @@ struct InstanceConfig: Equatable, Encodable {
     googleAI_v1beta_freeTier,
     // Note: The following configs are commented out for easy one-off manual testing.
     // vertexAI_v1beta_staging,
+    // vertexAI_v1beta_staging_global_bypassProxy,
     // googleAI_v1beta_staging,
     // googleAI_v1beta_freeTier_bypassProxy,
   ]
@@ -162,6 +169,8 @@ extension InstanceConfig: CustomTestStringConvertible {
       " - Staging"
     case .googleAIBypassProxy:
       " - Bypass Proxy"
+    case .vertexAIStagingBypassProxy:
+      " - Staging - Bypass Proxy"
     }
     let locationSuffix: String
     if case let .vertexAI(_, location: location) = apiConfig.service {