Skip to content

Commit e01df11

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: 852263160
1 parent b48b194 commit e01df11

File tree

11 files changed

+420
-203
lines changed

11 files changed

+420
-203
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: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
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+
/**
11+
* This interface provides a method to convert a service account JSON string to a Google Credentials
12+
* object.
13+
*
14+
* <p>Additionally, contains helper methods that aid with transfering the credentials' data to the
15+
* HttpRequest.Builder object
16+
*/
17+
public interface CredentialsHelper {
18+
19+
/**
20+
* Converts a service account JSON string to a Google Credentials object.
21+
*
22+
* @param serviceAccountJson The service account JSON string.
23+
* @return A Google Credentials object.
24+
* @throws IOException when an error occurs during the conversion.
25+
*/
26+
Credentials getGoogleCredentials(@Nullable String serviceAccountJson) throws IOException;
27+
28+
/**
29+
* Populates the headers (such as Authorization or x-goog-project) in the HttpRequest.Builder with
30+
* the metadata from the credentials.
31+
*
32+
* @param builder HttpRequest.Builder object to populate the headers
33+
* @param credentials Credentials object containing the metadata
34+
* @return HttpRequest.Builder object with the headers populated
35+
* @throws IOException if an error occurs when getting the metadata from the credentials
36+
*/
37+
public static HttpRequest.Builder populateHeaders(
38+
HttpRequest.Builder builder, Credentials credentials) throws IOException {
39+
for (Map.Entry<String, List<String>> entry : credentials.getRequestMetadata().entrySet()) {
40+
for (String value : entry.getValue()) {
41+
builder = builder.header(entry.getKey(), value);
42+
}
43+
}
44+
return builder;
45+
}
46+
}
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+
}

0 commit comments

Comments
 (0)