diff --git a/src/main/java/com/google/genai/gaos/hooks/SDKHooks.java b/src/main/java/com/google/genai/gaos/hooks/SDKHooks.java index 0755b4ae4289..dbd4dac50ca8 100644 --- a/src/main/java/com/google/genai/gaos/hooks/SDKHooks.java +++ b/src/main/java/com/google/genai/gaos/hooks/SDKHooks.java @@ -37,11 +37,12 @@ private SDKHooks() { public static void initialize(com.google.genai.gaos.utils.Hooks hooks) { hooks.registerBeforeRequest( (context, request) -> { + HttpRequest.Builder builder = null; if (context.securitySource().isPresent()) { HasSecurity hasSecurity = context.securitySource().get().getSecurity(); if (hasSecurity instanceof Security) { Security security = (Security) hasSecurity; - HttpRequest.Builder builder = request.toBuilder(); + builder = request.toBuilder(); if (security.defaultHeaders().isPresent()) { for (Map.Entry entry : security.defaultHeaders().get().entrySet()) { @@ -53,21 +54,33 @@ public static void initialize(com.google.genai.gaos.utils.Hooks hooks) { } else if (security.accessToken().isPresent()) { builder.setHeader("Authorization", "Bearer " + security.accessToken().get()); } - return builder.build(); } } - return request; + if (context.sdkConfiguration() != null + && context.sdkConfiguration().globals != null + && !request.headers().firstValue("x-goog-user-project").isPresent()) { + java.util.Optional userProject = + context.sdkConfiguration().globals.getParam("header", "x-goog-user-project"); + if (userProject.isPresent()) { + if (builder == null) { + builder = request.toBuilder(); + } + builder.setHeader("x-goog-user-project", userProject.get()); + } + } + return builder != null ? builder.build() : request; }); } public static void initialize(com.google.genai.gaos.utils.AsyncHooks asyncHooks) { asyncHooks.registerBeforeRequest( (context, request) -> { + HttpRequest.Builder builder = null; if (context.securitySource().isPresent()) { HasSecurity hasSecurity = context.securitySource().get().getSecurity(); if (hasSecurity instanceof Security) { Security security = (Security) hasSecurity; - HttpRequest.Builder builder = request.toBuilder(); + builder = request.toBuilder(); if (security.defaultHeaders().isPresent()) { for (Map.Entry entry : security.defaultHeaders().get().entrySet()) { @@ -79,10 +92,22 @@ public static void initialize(com.google.genai.gaos.utils.AsyncHooks asyncHooks) } else if (security.accessToken().isPresent()) { builder.setHeader("Authorization", "Bearer " + security.accessToken().get()); } - return CompletableFuture.completedFuture(builder.build()); } } - return CompletableFuture.completedFuture(request); + if (context.sdkConfiguration() != null + && context.sdkConfiguration().globals != null + && !request.headers().firstValue("x-goog-user-project").isPresent()) { + java.util.Optional userProject = + context.sdkConfiguration().globals.getParam("header", "x-goog-user-project"); + if (userProject.isPresent()) { + if (builder == null) { + builder = request.toBuilder(); + } + builder.setHeader("x-goog-user-project", userProject.get()); + } + } + return CompletableFuture.completedFuture( + builder != null ? builder.build() : request); }); } diff --git a/src/test/java/com/google/genai/gaos/hooks/SDKHooksTest.java b/src/test/java/com/google/genai/gaos/hooks/SdkHooksTest.java similarity index 72% rename from src/test/java/com/google/genai/gaos/hooks/SDKHooksTest.java rename to src/test/java/com/google/genai/gaos/hooks/SdkHooksTest.java index 9d7280094c15..39455f7ab392 100644 --- a/src/test/java/com/google/genai/gaos/hooks/SDKHooksTest.java +++ b/src/test/java/com/google/genai/gaos/hooks/SdkHooksTest.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; +import com.google.genai.gaos.SDKConfiguration; import com.google.genai.gaos.SecuritySource; import com.google.genai.gaos.models.shared.Security; import com.google.genai.gaos.utils.AsyncHooks; @@ -32,11 +33,11 @@ import org.junit.jupiter.api.Test; /** - * Tests for {@link SDKHooks} to verify that request authorization headers - * (x-goog-api-key, Authorization, defaultHeaders) are properly configured - * and not accidentally cleared during SDK regeneration. + * Tests for {@link SDKHooks} to verify that request authorization headers (x-goog-api-key, + * Authorization, defaultHeaders) are properly configured and not accidentally cleared during SDK + * regeneration. */ -public final class SDKHooksTest { +public final class SdkHooksTest { private static final String API_KEY = "test-api-key-12345"; private static final String ACCESS_TOKEN = "test-access-token-abcde"; @@ -80,7 +81,8 @@ public void testSyncHookSetsAuthorizationHeader() throws Exception { HttpRequest request = createBaseRequest(); HttpRequest modified = hooks.beforeRequest(context, request); - assertEquals("Bearer " + ACCESS_TOKEN, modified.headers().firstValue("Authorization").orElse(null)); + assertEquals( + "Bearer " + ACCESS_TOKEN, modified.headers().firstValue("Authorization").orElse(null)); } @Test @@ -124,7 +126,8 @@ public void testAsyncHookSetsAuthorizationHeader() throws Exception { HttpRequest request = createBaseRequest(); HttpRequest modified = asyncHooks.beforeRequest(context, request).get(); - assertEquals("Bearer " + ACCESS_TOKEN, modified.headers().firstValue("Authorization").orElse(null)); + assertEquals( + "Bearer " + ACCESS_TOKEN, modified.headers().firstValue("Authorization").orElse(null)); } @Test @@ -154,4 +157,41 @@ public void testHookWithoutSecurityPassesRequestThrough() throws Exception { HttpRequest modified = hooks.beforeRequest(context, request); assertSame(request, modified); } + + @Test + public void testSyncAndAsyncHooksPropagateUserProjectFromGlobals() throws Exception { + Hooks hooks = new Hooks(); + SDKHooks.initialize(hooks); + AsyncHooks asyncHooks = new AsyncHooks(); + SDKHooks.initialize(asyncHooks); + + SDKConfiguration config = new SDKConfiguration(); + config.globals.putParam("header", "x-goog-user-project", "my-quota-project"); + Hook.BeforeRequestContext context = + new Hook.BeforeRequestContextImpl( + config, + "https://generativelanguage.googleapis.com", + "createInteraction", + Optional.empty(), + Optional.empty()); + + HttpRequest request = createBaseRequest(); + HttpRequest syncModified = hooks.beforeRequest(context, request); + assertEquals( + "my-quota-project", syncModified.headers().firstValue("x-goog-user-project").orElse(null)); + + HttpRequest asyncModified = asyncHooks.beforeRequest(context, request).get(); + assertEquals( + "my-quota-project", asyncModified.headers().firstValue("x-goog-user-project").orElse(null)); + + // Existing x-goog-user-project header on the request should not be overwritten. + HttpRequest presetRequest = + request.toBuilder().setHeader("x-goog-user-project", "override-project").build(); + HttpRequest syncPreset = hooks.beforeRequest(context, presetRequest); + assertEquals( + "override-project", syncPreset.headers().firstValue("x-goog-user-project").orElse(null)); + HttpRequest asyncPreset = asyncHooks.beforeRequest(context, presetRequest).get(); + assertEquals( + "override-project", asyncPreset.headers().firstValue("x-goog-user-project").orElse(null)); + } }