Skip to content

Commit e22ce2a

Browse files
SunnyShu0925SunnyShu0925
authored
Merge branch 'main' into fix/llm-slug-truncate-utf16safe
2 parents f555622 + e5259fa commit e22ce2a

2 files changed

Lines changed: 127 additions & 7 deletions

File tree

‎extensions/google/transport-stream.test.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,11 +1133,40 @@ describe("google transport stream", () => {
11331133

11341134
expect(googleAuthMock).toHaveBeenCalledWith({
11351135
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
1136+
clientOptions: { transporterOptions: { timeout: 30_000 } },
11361137
});
11371138
expect(googleAuthGetAccessTokenMock).toHaveBeenCalledTimes(1);
11381139
expect(tokenFetchMock).not.toHaveBeenCalled();
11391140
});
11401141

1142+
it("bounds google-auth-library ADC token resolution at the Vertex owner", async () => {
1143+
const tempDir = await mkdtemp(
1144+
path.join(os.tmpdir(), "openclaw-google-vertex-authlib-timeout-"),
1145+
);
1146+
vi.stubEnv("GOOGLE_APPLICATION_CREDENTIALS", "");
1147+
vi.stubEnv("HOME", path.join(tempDir, "home"));
1148+
vi.stubEnv("APPDATA", "");
1149+
vi.useFakeTimers();
1150+
googleAuthGetAccessTokenMock
1151+
.mockReturnValueOnce(new Promise(() => {}))
1152+
.mockResolvedValueOnce("ya29.recovered-token");
1153+
1154+
const pendingRefresh = resolveGoogleVertexAuthorizedUserHeaders(vi.fn());
1155+
const refreshError = pendingRefresh.catch((error: unknown) => error);
1156+
await vi.waitFor(() => expect(googleAuthGetAccessTokenMock).toHaveBeenCalledOnce());
1157+
await vi.advanceTimersByTimeAsync(30_000);
1158+
1159+
await expect(refreshError).resolves.toMatchObject({
1160+
name: "TimeoutError",
1161+
message: "request timed out",
1162+
});
1163+
await expect(resolveGoogleVertexAuthorizedUserHeaders(vi.fn())).resolves.toEqual({
1164+
Authorization: "Bearer ya29.recovered-token",
1165+
});
1166+
expect(googleAuthMock).toHaveBeenCalledTimes(2);
1167+
expect(googleAuthGetAccessTokenMock).toHaveBeenCalledTimes(2);
1168+
});
1169+
11411170
it("does not cache google-auth ADC tokens when fallback expiry would exceed Date range", async () => {
11421171
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-authlib-expiry-"));
11431172
vi.useFakeTimers();
@@ -1314,6 +1343,56 @@ describe("google transport stream", () => {
13141343
expect(result.content).toEqual([{ type: "text", text: "ok" }]);
13151344
});
13161345

1346+
it("times out an authorized_user ADC token refresh", async () => {
1347+
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-adc-timeout-"));
1348+
const credentialsPath = path.join(tempDir, "application_default_credentials.json");
1349+
await writeFile(
1350+
credentialsPath,
1351+
JSON.stringify({
1352+
type: "authorized_user",
1353+
client_id: "client-id",
1354+
client_secret: "client-secret",
1355+
refresh_token: "timeout-refresh-token",
1356+
}),
1357+
"utf8",
1358+
);
1359+
vi.stubEnv("GOOGLE_APPLICATION_CREDENTIALS", credentialsPath);
1360+
vi.useFakeTimers();
1361+
1362+
let observedSignal: AbortSignal | undefined;
1363+
const tokenFetchMock = vi.fn((_input: string | URL | Request, init?: RequestInit) => {
1364+
const signal = init?.signal;
1365+
if (!signal) {
1366+
throw new Error("expected token refresh deadline signal");
1367+
}
1368+
observedSignal = signal;
1369+
const body = new ReadableStream<Uint8Array>({
1370+
start(controller) {
1371+
signal.addEventListener("abort", () => controller.error(signal.reason), { once: true });
1372+
},
1373+
});
1374+
return Promise.resolve(new Response(body, { status: 200 }));
1375+
});
1376+
1377+
const pendingRefresh = resolveGoogleVertexAuthorizedUserHeaders(tokenFetchMock);
1378+
// Attach the rejection handler before advancing fake time so the expected
1379+
// timeout cannot surface as an unhandled rejection between timer ticks.
1380+
const refreshError = pendingRefresh.catch((error: unknown) => error);
1381+
await vi.waitFor(() => expect(tokenFetchMock).toHaveBeenCalledOnce());
1382+
const signal = observedSignal;
1383+
if (!signal) {
1384+
throw new Error("expected token refresh deadline signal");
1385+
}
1386+
expect(signal.aborted).toBe(false);
1387+
await vi.advanceTimersByTimeAsync(30_000);
1388+
1389+
expect(signal.aborted).toBe(true);
1390+
await expect(refreshError).resolves.toMatchObject({
1391+
name: "TimeoutError",
1392+
message: "request timed out",
1393+
});
1394+
});
1395+
13171396
it("refreshes authorized_user ADC from a compressed token response", async () => {
13181397
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-adc-gzip-"));
13191398
const credentialsPath = path.join(tempDir, "application_default_credentials.json");

‎extensions/google/vertex-adc.ts‎

Lines changed: 48 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,15 @@ import { readFile } from "node:fs/promises";
44
import os from "node:os";
55
import path from "node:path";
66
import { gunzipSync } from "node:zlib";
7+
import { buildTimeoutAbortSignal } from "openclaw/plugin-sdk/extension-shared";
78
import {
89
asDateTimestampMs,
910
resolveExpiresAtMsFromDurationMs,
1011
resolveExpiresAtMsFromDurationSeconds,
1112
} from "openclaw/plugin-sdk/number-runtime";
1213
import { readResponseWithLimit } from "openclaw/plugin-sdk/response-limit-runtime";
1314
import { normalizeOptionalString } from "openclaw/plugin-sdk/string-coerce-runtime";
15+
import { withTimeout } from "openclaw/plugin-sdk/text-utility-runtime";
1416

1517
type GoogleAuthorizedUserCredentials = {
1618
type: "authorized_user";
@@ -41,6 +43,7 @@ type GoogleOauthTokenResponsePayload = {
4143
const GCP_VERTEX_CREDENTIALS_MARKER = "gcp-vertex-credentials";
4244
const GOOGLE_OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token";
4345
const GOOGLE_VERTEX_OAUTH_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
46+
const GOOGLE_VERTEX_ADC_TOKEN_REFRESH_TIMEOUT_MS = 30_000;
4447
// Hold tokens slightly less long than reported expiry (Google's recommendation
4548
// is a 60s buffer) so we don't ship a request that's already revoked when it
4649
// leaves the gateway.
@@ -243,12 +246,26 @@ async function refreshGoogleVertexAuthorizedUserAccessToken(params: {
243246
refresh_token: refreshToken,
244247
grant_type: "refresh_token",
245248
});
246-
const response = await (params.fetchImpl ?? fetch)(GOOGLE_OAUTH_TOKEN_URL, {
247-
method: "POST",
248-
headers: { "Content-Type": "application/x-www-form-urlencoded" },
249-
body,
249+
const { signal, cleanup } = buildTimeoutAbortSignal({
250+
timeoutMs: GOOGLE_VERTEX_ADC_TOKEN_REFRESH_TIMEOUT_MS,
251+
operation: "google-vertex-adc-token-refresh",
252+
url: GOOGLE_OAUTH_TOKEN_URL,
250253
});
251-
const payload = await readGoogleOauthTokenResponsePayload(response);
254+
let response: Response;
255+
let payload: GoogleOauthTokenResponsePayload | undefined;
256+
try {
257+
response = await (params.fetchImpl ?? fetch)(GOOGLE_OAUTH_TOKEN_URL, {
258+
method: "POST",
259+
headers: { "Content-Type": "application/x-www-form-urlencoded" },
260+
body,
261+
signal,
262+
});
263+
// Keep the request deadline active through body consumption. Fetch resolves
264+
// at headers, so cleanup here would leave a stalled token body unbounded.
265+
payload = await readGoogleOauthTokenResponsePayload(response);
266+
} finally {
267+
cleanup();
268+
}
252269
if (!response.ok) {
253270
const description = normalizeOptionalString(payload?.error_description);
254271
const code = normalizeOptionalString(payload?.error);
@@ -345,18 +362,42 @@ async function resolveGoogleVertexAccessTokenViaGoogleAuth(): Promise<string> {
345362
// It also caches tokens internally and refreshes before expiry.
346363
return new GoogleAuth({
347364
scopes: [GOOGLE_VERTEX_OAUTH_SCOPE],
365+
// Best-effort cancellation for clients that use the shared transporter.
366+
// WIF STS and GCE metadata need the owner-level deadline below.
367+
clientOptions: {
368+
transporterOptions: { timeout: GOOGLE_VERTEX_ADC_TOKEN_REFRESH_TIMEOUT_MS },
369+
},
348370
});
349371
}),
350372
};
351373
}
352-
const auth = await cachedGoogleAuthClient.promise;
374+
const authClient = cachedGoogleAuthClient;
375+
const auth = await authClient.promise;
353376

354377
const cached = cachedGoogleVertexAdcToken;
355378
if (cached && isGoogleVertexTokenFresh(cached.expiresAtMs)) {
356379
return cached.token;
357380
}
358381

359-
const token = await auth.getAccessToken();
382+
// Some google-auth-library ADC implementations bypass the configured Gaxios
383+
// transporter, so this owner-level deadline also bounds STS and metadata paths.
384+
let token: string | null | undefined;
385+
try {
386+
token = await withTimeout(auth.getAccessToken(), GOOGLE_VERTEX_ADC_TOKEN_REFRESH_TIMEOUT_MS, {
387+
createError: () => new DOMException("request timed out", "TimeoutError"),
388+
});
389+
} catch (error) {
390+
// The dependency coalesces in-flight refreshes. Drop only this timed-out
391+
// client so a recovered identity endpoint gets a fresh attempt next time.
392+
if (
393+
error instanceof DOMException &&
394+
error.name === "TimeoutError" &&
395+
cachedGoogleAuthClient === authClient
396+
) {
397+
cachedGoogleAuthClient = undefined;
398+
}
399+
throw error;
400+
}
360401
const normalized = normalizeOptionalString(token);
361402
if (!normalized) {
362403
throw new Error(

0 commit comments

Comments
 (0)