Skip to content

Commit 03891e9

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: use Credentials' request metadata to populate headers
Instead of manually adding an authorization header and potentially omitting other necessary metadata (such as x-goog-user-project) copy all the necessary headers managed by the Credentials object into the request. PiperOrigin-RevId: 846816800
1 parent 9611f89 commit 03891e9

File tree

11 files changed

+343
-210
lines changed

11 files changed

+343
-210
lines changed

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/ApplicationIntegrationToolset.java

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,8 @@
88
import com.google.adk.agents.ReadonlyContext;
99
import com.google.adk.tools.BaseTool;
1010
import com.google.adk.tools.BaseToolset;
11-
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.DefaultHttpExecutor;
12-
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
1311
import io.reactivex.rxjava3.core.Flowable;
12+
import java.net.http.HttpClient;
1413
import java.util.ArrayList;
1514
import java.util.Iterator;
1615
import java.util.List;
@@ -30,7 +29,8 @@ public class ApplicationIntegrationToolset implements BaseToolset {
3029
@Nullable String toolNamePrefix;
3130
@Nullable String toolInstructions;
3231
public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
33-
HttpExecutor httpExecutor;
32+
private final HttpClient httpClient;
33+
private final CredentialsHelper credentialsHelper;
3434

3535
/**
3636
* ApplicationIntegrationToolset generates tools from a given Application Integration resource.
@@ -85,7 +85,8 @@ public ApplicationIntegrationToolset(
8585
serviceAccountJson,
8686
toolNamePrefix,
8787
toolInstructions,
88-
new DefaultHttpExecutor().createExecutor(serviceAccountJson));
88+
HttpClient.newHttpClient(),
89+
new GoogleCredentialsHelper());
8990
}
9091

9192
ApplicationIntegrationToolset(
@@ -99,7 +100,8 @@ public ApplicationIntegrationToolset(
99100
String serviceAccountJson,
100101
String toolNamePrefix,
101102
String toolInstructions,
102-
HttpExecutor httpExecutor) {
103+
HttpClient httpClient,
104+
CredentialsHelper credentialsHelper) {
103105
this.project = project;
104106
this.location = location;
105107
this.integration = integration;
@@ -110,7 +112,8 @@ public ApplicationIntegrationToolset(
110112
this.serviceAccountJson = serviceAccountJson;
111113
this.toolNamePrefix = toolNamePrefix;
112114
this.toolInstructions = toolInstructions;
113-
this.httpExecutor = httpExecutor;
115+
this.httpClient = httpClient;
116+
this.credentialsHelper = credentialsHelper;
114117
}
115118

116119
List<String> getPathUrl(String openApiSchemaString) throws Exception {
@@ -145,7 +148,9 @@ private List<BaseTool> getAllTools() throws Exception {
145148
null,
146149
null,
147150
null,
148-
this.httpExecutor);
151+
this.serviceAccountJson,
152+
this.httpClient,
153+
this.credentialsHelper);
149154
openApiSchemaString = integrationClient.generateOpenApiSpec();
150155
List<String> pathUrls = getPathUrl(openApiSchemaString);
151156
for (String pathUrl : pathUrls) {
@@ -161,7 +166,8 @@ private List<BaseTool> getAllTools() throws Exception {
161166
null,
162167
null,
163168
this.serviceAccountJson,
164-
this.httpExecutor));
169+
this.httpClient,
170+
this.credentialsHelper));
165171
}
166172
}
167173
} else if (!isNullOrEmpty(this.connection)
@@ -175,7 +181,9 @@ private List<BaseTool> getAllTools() throws Exception {
175181
this.connection,
176182
this.entityOperations,
177183
this.actions,
178-
this.httpExecutor);
184+
this.serviceAccountJson,
185+
this.httpClient,
186+
this.credentialsHelper);
179187
ObjectNode parentOpenApiSpec = OBJECT_MAPPER.createObjectNode();
180188
ObjectNode openApiSpec =
181189
integrationClient.getOpenApiSpecForConnection(toolNamePrefix, toolInstructions);
@@ -188,7 +196,14 @@ private List<BaseTool> getAllTools() throws Exception {
188196
if (!isNullOrEmpty(toolName)) {
189197
ConnectionsClient connectionsClient =
190198
new ConnectionsClient(
191-
this.project, this.location, this.connection, this.httpExecutor, OBJECT_MAPPER);
199+
this.project,
200+
this.location,
201+
this.connection,
202+
this.serviceAccountJson,
203+
this.httpClient,
204+
this.credentialsHelper,
205+
OBJECT_MAPPER);
206+
192207
ConnectionsClient.ConnectionDetails connectionDetails =
193208
connectionsClient.getConnectionDetails();
194209

@@ -202,7 +217,8 @@ private List<BaseTool> getAllTools() throws Exception {
202217
connectionDetails.serviceName,
203218
connectionDetails.host,
204219
this.serviceAccountJson,
205-
this.httpExecutor));
220+
this.httpClient,
221+
this.credentialsHelper));
206222
}
207223
}
208224
} else {

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/ConnectionsClient.java

Lines changed: 35 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44

55
import com.fasterxml.jackson.core.type.TypeReference;
66
import com.fasterxml.jackson.databind.ObjectMapper;
7-
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
7+
import com.google.auth.Credentials;
8+
import com.google.common.base.Preconditions;
89
import com.google.common.collect.ImmutableList;
910
import com.google.common.collect.ImmutableMap;
1011
import java.io.IOException;
1112
import java.net.URI;
13+
import java.net.http.HttpClient;
1214
import java.net.http.HttpRequest;
1315
import java.net.http.HttpResponse;
1416
import java.util.HashMap;
@@ -28,7 +30,9 @@ public class ConnectionsClient {
2830
private final String location;
2931
private final String connection;
3032
private static final String CONNECTOR_URL = "https://connectors.googleapis.com";
31-
private final HttpExecutor httpExecutor;
33+
private final HttpClient httpClient;
34+
private final String serviceAccountJson;
35+
private final CredentialsHelper credentialsHelper;
3236
private final ObjectMapper objectMapper;
3337

3438
/** Represents details of a connection. */
@@ -63,13 +67,34 @@ public ConnectionsClient(
6367
String project,
6468
String location,
6569
String connection,
66-
HttpExecutor httpExecutor,
70+
String serviceAccountJson,
71+
HttpClient httpClient,
72+
CredentialsHelper credentialsHelper,
6773
ObjectMapper objectMapper) {
74+
6875
this.project = project;
6976
this.location = location;
7077
this.connection = connection;
71-
this.httpExecutor = httpExecutor;
78+
this.httpClient = Preconditions.checkNotNull(httpClient);
7279
this.objectMapper = objectMapper;
80+
this.serviceAccountJson = serviceAccountJson;
81+
this.credentialsHelper = Preconditions.checkNotNull(credentialsHelper);
82+
}
83+
84+
public ConnectionsClient(
85+
String project,
86+
String location,
87+
String connection,
88+
HttpClient httpClient,
89+
ObjectMapper objectMapper) {
90+
this(
91+
project,
92+
location,
93+
connection,
94+
null,
95+
httpClient,
96+
new GoogleCredentialsHelper(),
97+
objectMapper);
7398
}
7499

75100
/**
@@ -173,16 +198,17 @@ public ActionSchema getActionSchema(String action) throws IOException, Interrupt
173198
}
174199

175200
private HttpResponse<String> executeApiCall(String url) throws IOException, InterruptedException {
176-
HttpRequest request =
201+
HttpRequest.Builder requestBuilder =
177202
HttpRequest.newBuilder()
178203
.uri(URI.create(url))
179204
.header("Content-Type", "application/json")
180-
.header("Authorization", "Bearer " + httpExecutor.getToken())
181-
.GET()
182-
.build();
205+
.GET();
206+
207+
Credentials credentials = credentialsHelper.getGoogleCredentials(serviceAccountJson);
208+
requestBuilder = CredentialsHelper.populateHeaders(requestBuilder, credentials);
183209

184210
HttpResponse<String> response =
185-
httpExecutor.send(request, HttpResponse.BodyHandlers.ofString());
211+
httpClient.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofString());
186212

187213
if (response.statusCode() >= 400) {
188214
String body = response.body();
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package com.google.adk.tools.applicationintegrationtoolset;
2+
3+
import com.google.auth.Credentials;
4+
import java.io.IOException;
5+
import java.net.http.HttpRequest;
6+
import java.util.List;
7+
import java.util.Map;
8+
import javax.annotation.Nullable;
9+
10+
public interface CredentialsHelper {
11+
12+
Credentials getGoogleCredentials(@Nullable String serviceAccountJson) throws IOException;
13+
14+
public static HttpRequest.Builder populateHeaders(
15+
HttpRequest.Builder builder, Credentials credentials) throws IOException {
16+
for (Map.Entry<String, List<String>> entry : credentials.getRequestMetadata().entrySet()) {
17+
for (String value : entry.getValue()) {
18+
builder = builder.header(entry.getKey(), value);
19+
}
20+
}
21+
return builder;
22+
}
23+
}
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package com.google.adk.tools.applicationintegrationtoolset;
2+
3+
import static java.nio.charset.StandardCharsets.UTF_8;
4+
5+
import com.google.auth.oauth2.GoogleCredentials;
6+
import com.google.auth.oauth2.ServiceAccountCredentials;
7+
import java.io.ByteArrayInputStream;
8+
import java.io.IOException;
9+
import java.io.InputStream;
10+
import javax.annotation.Nullable;
11+
12+
public final class GoogleCredentialsHelper implements CredentialsHelper {
13+
14+
@Override
15+
public GoogleCredentials getGoogleCredentials(@Nullable String serviceAccountJson)
16+
throws IOException {
17+
GoogleCredentials credentials;
18+
19+
if (serviceAccountJson != null && !serviceAccountJson.isBlank()) {
20+
try (InputStream is = new ByteArrayInputStream(serviceAccountJson.getBytes(UTF_8))) {
21+
credentials = ServiceAccountCredentials.fromStream(is);
22+
}
23+
} else {
24+
credentials = GoogleCredentials.getApplicationDefault();
25+
}
26+
credentials = credentials.createScoped("https://www.googleapis.com/auth/cloud-platform");
27+
credentials.refreshIfExpired();
28+
return credentials;
29+
}
30+
}

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/IntegrationClient.java

Lines changed: 75 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55
import com.fasterxml.jackson.databind.JsonNode;
66
import com.fasterxml.jackson.databind.ObjectMapper;
77
import com.fasterxml.jackson.databind.node.ObjectNode;
8-
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
8+
import com.google.auth.Credentials;
99
import com.google.common.base.Preconditions;
1010
import com.google.common.collect.ImmutableList;
1111
import com.google.common.collect.ImmutableMap;
1212
import java.io.IOException;
1313
import java.net.URI;
14+
import java.net.http.HttpClient;
1415
import java.net.http.HttpRequest;
1516
import java.net.http.HttpResponse;
1617
import java.util.Arrays;
@@ -26,14 +27,17 @@
2627
* <p>This class provides methods for retrieving OpenAPI spec for an integration or a connection.
2728
*/
2829
public class IntegrationClient {
29-
String project;
30-
String location;
31-
String integration;
32-
List<String> triggers;
33-
String connection;
34-
Map<String, List<String>> entityOperations;
35-
List<String> actions;
36-
private final HttpExecutor httpExecutor;
30+
private final String project;
31+
private final String location;
32+
private final String integration;
33+
private final List<String> triggers;
34+
private final String connection;
35+
private final Map<String, List<String>> entityOperations;
36+
private final List<String> actions;
37+
private final String serviceAccountJson;
38+
private final HttpClient httpClient;
39+
private final CredentialsHelper credentialsHelper;
40+
3741
public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
3842

3943
IntegrationClient(
@@ -44,20 +48,66 @@ public class IntegrationClient {
4448
String connection,
4549
Map<String, List<String>> entityOperations,
4650
List<String> actions,
47-
HttpExecutor httpExecutor) {
51+
String serviceAccountJson,
52+
HttpClient httpClient,
53+
CredentialsHelper credentialsHelper) {
54+
4855
this.project = project;
4956
this.location = location;
5057
this.integration = integration;
5158
this.triggers = triggers;
5259
this.connection = connection;
5360
this.entityOperations = entityOperations;
5461
this.actions = actions;
55-
this.httpExecutor = httpExecutor;
62+
this.serviceAccountJson = serviceAccountJson;
63+
this.httpClient = Preconditions.checkNotNull(httpClient);
64+
this.credentialsHelper = Preconditions.checkNotNull(credentialsHelper);
5665
if (!isNullOrEmpty(connection)) {
5766
validate();
5867
}
5968
}
6069

70+
IntegrationClient(
71+
String project,
72+
String location,
73+
String integration,
74+
List<String> triggers,
75+
String connection,
76+
Map<String, List<String>> entityOperations,
77+
List<String> actions) {
78+
this(
79+
project,
80+
location,
81+
integration,
82+
triggers,
83+
connection,
84+
entityOperations,
85+
actions,
86+
HttpClient.newHttpClient());
87+
}
88+
89+
IntegrationClient(
90+
String project,
91+
String location,
92+
String integration,
93+
List<String> triggers,
94+
String connection,
95+
Map<String, List<String>> entityOperations,
96+
List<String> actions,
97+
HttpClient httpClient) {
98+
this(
99+
project,
100+
location,
101+
integration,
102+
triggers,
103+
connection,
104+
entityOperations,
105+
actions,
106+
null,
107+
httpClient,
108+
new GoogleCredentialsHelper());
109+
}
110+
61111
private void validate() {
62112
// Check if both are null, throw exception
63113

@@ -116,15 +166,17 @@ String generateOpenApiSpec() throws Exception {
116166
Arrays.asList(this.triggers))),
117167
"fileFormat",
118168
"JSON"));
119-
HttpRequest request =
169+
HttpRequest.Builder requestBuilder =
120170
HttpRequest.newBuilder()
121171
.uri(URI.create(url))
122-
.header("Authorization", "Bearer " + httpExecutor.getToken())
123172
.header("Content-Type", "application/json")
124-
.POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody))
125-
.build();
173+
.POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody));
174+
175+
Credentials credentials = credentialsHelper.getGoogleCredentials(serviceAccountJson);
176+
requestBuilder = CredentialsHelper.populateHeaders(requestBuilder, credentials);
177+
126178
HttpResponse<String> response =
127-
httpExecutor.send(request, HttpResponse.BodyHandlers.ofString());
179+
httpClient.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofString());
128180

129181
if (response.statusCode() < 200 || response.statusCode() >= 300) {
130182
throw new Exception("Error fetching OpenAPI spec. Status: " + response.statusCode());
@@ -326,6 +378,12 @@ String getOperationIdFromPathUrl(String openApiSchemaString, String pathUrl) thr
326378

327379
ConnectionsClient createConnectionsClient() {
328380
return new ConnectionsClient(
329-
this.project, this.location, this.connection, this.httpExecutor, OBJECT_MAPPER);
381+
this.project,
382+
this.location,
383+
this.connection,
384+
this.serviceAccountJson,
385+
this.httpClient,
386+
this.credentialsHelper,
387+
OBJECT_MAPPER);
330388
}
331389
}

0 commit comments

Comments
 (0)