From 33a44a979127a4c2b5b19e7228d65939aa55f967 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Thu, 31 Oct 2024 17:59:04 -0400
Subject: [PATCH] [Vertex AI] Fix caching of Vertex AI instances (#14007)

---
 FirebaseVertexAI/CHANGELOG.md                 |  5 ++
 FirebaseVertexAI/Sources/VertexAI.swift       |  9 +--
 .../Tests/Unit/VertexComponentTests.swift     | 67 +++++++++++++------
 3 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/FirebaseVertexAI/CHANGELOG.md b/FirebaseVertexAI/CHANGELOG.md
index cd290c5c1be..870f0c53788 100644
--- a/FirebaseVertexAI/CHANGELOG.md
+++ b/FirebaseVertexAI/CHANGELOG.md
@@ -1,3 +1,8 @@
+# 11.5.0
+- [fixed] Fixed an issue where `VertexAI.vertexAI(app: app1)` and
+  `VertexAI.vertexAI(app: app2)` would return the same instance if their
+  `location` was the same, including the default `us-central1`. (#14007)
+
 # 11.4.0
 - [feature] Vertex AI in Firebase is now Generally Available (GA) and can be
   used in production apps. (#13725)
diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift
index ad378c067d1..c0cd2cb66a3 100644
--- a/FirebaseVertexAI/Sources/VertexAI.swift
+++ b/FirebaseVertexAI/Sources/VertexAI.swift
@@ -55,11 +55,12 @@ public class VertexAI {
     // Unlock before the function returns.
     defer { os_unfair_lock_unlock(&instancesLock) }
 
-    if let instance = instances[location] {
+    let instanceKey = "\(app.name):\(location)"
+    if let instance = instances[instanceKey] {
       return instance
     }
     let newInstance = VertexAI(app: app, location: location)
-    instances[location] = newInstance
+    instances[instanceKey] = newInstance
     return newInstance
   }
 
@@ -116,8 +117,8 @@ public class VertexAI {
 
   private let auth: AuthInterop?
 
-  /// A map of active  `VertexAI` instances for `app`, keyed by model resource names
-  /// (e.g., "projects/my-project-id/locations/us-central1/publishers/google/models/gemini-pro").
+  /// A map of active  `VertexAI` instances keyed by the `FirebaseApp` name and the `location`, in
+  /// the format `appName:location`.
   private static var instances: [String: VertexAI] = [:]
 
   /// Lock to manage access to the `instances` array to avoid race conditions.
diff --git a/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift b/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift
index bea85bf96e7..5e685dd98bd 100644
--- a/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift
+++ b/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift
@@ -24,22 +24,21 @@ import XCTest
 class VertexComponentTests: XCTestCase {
   static let projectID = "test-project-id"
   static let apiKey = "test-api-key"
+  static let options = {
+    let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
+                                  gcmSenderID: "00000000000000000-00000000000-000000000")
+    options.projectID = VertexComponentTests.projectID
+    options.apiKey = VertexComponentTests.apiKey
 
-  static var app: FirebaseApp?
+    return options
+  }()
 
-  let location = "test-location"
+  static let app = {
+    FirebaseApp.configure(options: options)
+    return FirebaseApp(instanceWithName: "test", options: options)
+  }()
 
-  override class func setUp() {
-    super.setUp()
-    if app == nil {
-      let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
-                                    gcmSenderID: "00000000000000000-00000000000-000000000")
-      options.projectID = VertexComponentTests.projectID
-      options.apiKey = VertexComponentTests.apiKey
-      FirebaseApp.configure(options: options)
-      app = FirebaseApp(instanceWithName: "test", options: options)
-    }
-  }
+  let location = "test-location"
 
   /// Test that the objc class is available for the component system to update the user agent.
   func testComponentsBeingRegistered() throws {
@@ -48,9 +47,7 @@ class VertexComponentTests: XCTestCase {
 
   /// Tests that a vertex instance can be created properly.
   func testVertexInstanceCreation() throws {
-    let app = try XCTUnwrap(VertexComponentTests.app)
-
-    let vertex = VertexAI.vertexAI(app: app, location: location)
+    let vertex = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
 
     XCTAssertNotNil(vertex)
     XCTAssertEqual(vertex.projectID, VertexComponentTests.projectID)
@@ -58,17 +55,47 @@ class VertexComponentTests: XCTestCase {
     XCTAssertEqual(vertex.location, location)
   }
 
-  /// Tests that a vertex instances are reused properly.
-  func testMultipleComponentInstancesCreated() throws {
+  /// Tests that Vertex instances are reused properly.
+  func testSameAppAndLocation_instanceReused() throws {
     let app = try XCTUnwrap(VertexComponentTests.app)
+
     let vertex1 = VertexAI.vertexAI(app: app, location: location)
     let vertex2 = VertexAI.vertexAI(app: app, location: location)
 
     // Ensure they're the same instance.
     XCTAssert(vertex1 === vertex2)
+  }
+
+  func testSameAppAndDifferentLocation_newInstanceCreated() throws {
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: VertexComponentTests.app, location: "differentLocation")
+
+    // Ensure they are different instances.
+    XCTAssert(vertex1 !== vertex2)
+  }
+
+  func testDifferentAppAndSameLocation_newInstanceCreated() throws {
+    FirebaseApp.configure(name: "test-2", options: VertexComponentTests.options)
+    let app2 = FirebaseApp(instanceWithName: "test-2", options: VertexComponentTests.options)
+    addTeardownBlock { await app2.delete() }
+
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: app2, location: location)
+
+    XCTAssert(VertexComponentTests.app != app2)
+    XCTAssert(vertex1 !== vertex2) // Ensure they are different instances.
+  }
+
+  func testDifferentAppAndDifferentLocation_newInstanceCreated() throws {
+    FirebaseApp.configure(name: "test-2", options: VertexComponentTests.options)
+    let app2 = FirebaseApp(instanceWithName: "test-2", options: VertexComponentTests.options)
+    addTeardownBlock { await app2.delete() }
+
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: app2, location: "differentLocation")
 
-    let vertex3 = VertexAI.vertexAI(app: app, location: "differentLocation")
-    XCTAssert(vertex1 !== vertex3)
+    XCTAssert(VertexComponentTests.app != app2)
+    XCTAssert(vertex1 !== vertex2) // Ensure they are different instances.
   }
 
   /// Test that vertex instances get deallocated.