fix: customHeaders are not passed to the VertexAI clients (#11506)

PiperOrigin-RevId: 738911474

Co-authored-by: Jaycee Li <[email protected]>

Author: copybara-service[bot]

java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java

@@ -360,13 +360,6 @@ private PredictionServiceSettings getPredictionServiceSettings() throws IOExcept
     builder.setEndpoint(String.format("%s:443", apiEndpoint));
     builder.setCredentialsProvider(credentialsProvider);
 
-    HeaderProvider headerProvider =
-        FixedHeaderProvider.create(
-            "user-agent",
-            String.format(
-                "%s/%s",
-                Constants.USER_AGENT_HEADER,
-                GaxProperties.getLibraryVersion(PredictionServiceSettings.class)));
     builder.setHeaderProvider(headerProvider);
     return builder.build();
   }
@@ -435,13 +428,6 @@ private LlmUtilityServiceSettings getLlmUtilityServiceClientSettings() throws IO
     settingsBuilder.setEndpoint(String.format("%s:443", apiEndpoint));
     settingsBuilder.setCredentialsProvider(credentialsProvider);
 
-    HeaderProvider headerProvider =
-        FixedHeaderProvider.create(
-            "user-agent",
-            String.format(
-                "%s/%s",
-                Constants.USER_AGENT_HEADER,
-                GaxProperties.getLibraryVersion(LlmUtilityServiceSettings.class)));
     settingsBuilder.setHeaderProvider(headerProvider);
     return settingsBuilder.build();
   }

java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java

@@ -25,6 +25,8 @@
 import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.core.GoogleCredentialsProvider;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.vertexai.api.LlmUtilityServiceClient;
+import com.google.cloud.vertexai.api.LlmUtilityServiceSettings;
 import com.google.cloud.vertexai.api.PredictionServiceClient;
 import com.google.cloud.vertexai.api.PredictionServiceSettings;
 import com.google.common.collect.ImmutableList;
@@ -58,6 +60,8 @@ public final class VertexAITest {
 
   @Mock private PredictionServiceClient mockPredictionServiceClient;
 
+  @Mock private LlmUtilityServiceClient mockLlmUtilityServiceClient;
+
   @Mock private GoogleCredentialsProvider.Builder mockCredentialsProviderBuilder;
 
   @Mock private GoogleCredentialsProvider mockCredentialsProvider;
@@ -425,6 +429,20 @@ public void testInstantiateVertexAI_builderWithCustomHeaders_shouldContainRightF
             Constants.USER_AGENT_HEADER,
             GaxProperties.getLibraryVersion(PredictionServiceSettings.class)));
     assertThat(vertexAi.getHeaders()).isEqualTo(expectedHeaders);
+
+    // make sure the custom headers are set correctly in the prediction service client
+    try (MockedStatic mockStatic = mockStatic(PredictionServiceClient.class)) {
+      mockStatic
+          .when(() -> PredictionServiceClient.create(any(PredictionServiceSettings.class)))
+          .thenReturn(mockPredictionServiceClient);
+      PredictionServiceClient unused = vertexAi.getPredictionServiceClient();
+
+      ArgumentCaptor<PredictionServiceSettings> settings =
+          ArgumentCaptor.forClass(PredictionServiceSettings.class);
+      mockStatic.verify(() -> PredictionServiceClient.create(settings.capture()));
+
+      assertThat(settings.getValue().getHeaderProvider().getHeaders()).isEqualTo(expectedHeaders);
+    }
   }
 
   @Test
@@ -454,5 +472,19 @@ public void testInstantiateVertexAI_builderWithCustomHeaders_shouldContainRightF
             GaxProperties.getLibraryVersion(PredictionServiceSettings.class),
             "test_value"));
     assertThat(vertexAi.getHeaders()).isEqualTo(expectedHeaders);
+
+    // make sure the custom headers are set correctly in the llm utility service client
+    try (MockedStatic mockStatic = mockStatic(LlmUtilityServiceClient.class)) {
+      mockStatic
+          .when(() -> LlmUtilityServiceClient.create(any(LlmUtilityServiceSettings.class)))
+          .thenReturn(mockLlmUtilityServiceClient);
+      LlmUtilityServiceClient unused = vertexAi.getLlmUtilityClient();
+
+      ArgumentCaptor<LlmUtilityServiceSettings> settings =
+          ArgumentCaptor.forClass(LlmUtilityServiceSettings.class);
+      mockStatic.verify(() -> LlmUtilityServiceClient.create(settings.capture()));
+
+      assertThat(settings.getValue().getHeaderProvider().getHeaders()).isEqualTo(expectedHeaders);
+    }
   }
 }