Skip to content

Commit 39c2079

Browse files
damianFelixPagocursoragent
damianFelixPago
and
cursoragent
committed
fix(google-vertex): support production ADC modes
Co-authored-by: Cursor <[email protected]>
1 parent 6fcfeed commit 39c2079

8 files changed

Lines changed: 316 additions & 33 deletions

File tree

‎extensions/google/package.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
"type": "module",
77
"dependencies": {
88
"@earendil-works/pi-ai": "0.75.1",
9-
"@google/genai": "2.2.0"
9+
"@google/genai": "2.2.0",
10+
"google-auth-library": "10.6.2"
1011
},
1112
"devDependencies": {
1213
"@openclaw/plugin-sdk": "workspace:*"

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

Lines changed: 120 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,40 @@ import path from "node:path";
44
import type { Model } from "@earendil-works/pi-ai";
55
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
66

7-
const { buildGuardedModelFetchMock, guardedFetchMock } = vi.hoisted(() => ({
8-
buildGuardedModelFetchMock: vi.fn(),
9-
guardedFetchMock: vi.fn(),
10-
}));
7+
const {
8+
buildGuardedModelFetchMock,
9+
guardedFetchMock,
10+
googleAuthGetAccessTokenMock,
11+
googleAuthMock,
12+
} = vi.hoisted(() => {
13+
const googleAuthGetAccessTokenMock = vi.fn();
14+
return {
15+
buildGuardedModelFetchMock: vi.fn(),
16+
guardedFetchMock: vi.fn(),
17+
googleAuthGetAccessTokenMock,
18+
googleAuthMock: vi.fn(function GoogleAuthMock() {
19+
return {
20+
getAccessToken: googleAuthGetAccessTokenMock,
21+
};
22+
}),
23+
};
24+
});
1125

1226
vi.mock("openclaw/plugin-sdk/provider-transport-runtime", async (importOriginal) => ({
1327
...(await importOriginal()),
1428
buildGuardedModelFetch: buildGuardedModelFetchMock,
1529
}));
1630

31+
vi.mock("google-auth-library", () => ({
32+
GoogleAuth: googleAuthMock,
33+
}));
34+
1735
let buildGoogleGenerativeAiParams: typeof import("./transport-stream.js").buildGoogleGenerativeAiParams;
1836
let buildGoogleGemini3FirstResponseRetryParams: typeof import("./transport-stream.js").buildGoogleGemini3FirstResponseRetryParams;
1937
let createGoogleGenerativeAiTransportStreamFn: typeof import("./transport-stream.js").createGoogleGenerativeAiTransportStreamFn;
2038
let createGoogleVertexTransportStreamFn: typeof import("./transport-stream.js").createGoogleVertexTransportStreamFn;
2139
let hasGoogleVertexAuthorizedUserAdcSync: typeof import("./vertex-adc.js").hasGoogleVertexAuthorizedUserAdcSync;
40+
let resolveGoogleVertexAuthorizedUserHeaders: typeof import("./vertex-adc.js").resolveGoogleVertexAuthorizedUserHeaders;
2241
let resetGoogleVertexAuthorizedUserTokenCacheForTest: typeof import("./vertex-adc.js").resetGoogleVertexAuthorizedUserTokenCacheForTest;
2342

2443
const MODEL_PROVIDER_REQUEST_TRANSPORT_SYMBOL = Symbol.for(
@@ -254,13 +273,18 @@ describe("google transport stream", () => {
254273
createGoogleGenerativeAiTransportStreamFn,
255274
createGoogleVertexTransportStreamFn,
256275
} = await import("./transport-stream.js"));
257-
({ hasGoogleVertexAuthorizedUserAdcSync, resetGoogleVertexAuthorizedUserTokenCacheForTest } =
258-
await import("./vertex-adc.js"));
276+
({
277+
hasGoogleVertexAuthorizedUserAdcSync,
278+
resolveGoogleVertexAuthorizedUserHeaders,
279+
resetGoogleVertexAuthorizedUserTokenCacheForTest,
280+
} = await import("./vertex-adc.js"));
259281
});
260282

261283
beforeEach(() => {
262284
buildGuardedModelFetchMock.mockReset();
263285
guardedFetchMock.mockReset();
286+
googleAuthGetAccessTokenMock.mockReset();
287+
googleAuthMock.mockClear();
264288
buildGuardedModelFetchMock.mockReturnValue(guardedFetchMock);
265289
resetGoogleVertexAuthorizedUserTokenCacheForTest();
266290
});
@@ -271,6 +295,7 @@ describe("google transport stream", () => {
271295

272296
afterAll(() => {
273297
vi.doUnmock("openclaw/plugin-sdk/provider-transport-runtime");
298+
vi.doUnmock("google-auth-library");
274299
vi.resetModules();
275300
});
276301

@@ -695,6 +720,95 @@ describe("google transport stream", () => {
695720
});
696721
});
697722

723+
it("detects supported Vertex ADC sources synchronously", async () => {
724+
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-adc-detect-"));
725+
for (const type of ["authorized_user", "external_account", "service_account"]) {
726+
const credentialsPath = path.join(tempDir, `${type}.json`);
727+
await writeFile(credentialsPath, JSON.stringify({ type }), "utf8");
728+
729+
expect(
730+
hasGoogleVertexAuthorizedUserAdcSync({
731+
GOOGLE_APPLICATION_CREDENTIALS: credentialsPath,
732+
}),
733+
).toBe(true);
734+
}
735+
736+
expect(
737+
hasGoogleVertexAuthorizedUserAdcSync({
738+
HOME: path.join(tempDir, "empty-home"),
739+
KUBERNETES_SERVICE_HOST: "10.0.0.1",
740+
}),
741+
).toBe(true);
742+
expect(
743+
hasGoogleVertexAuthorizedUserAdcSync({
744+
HOME: path.join(tempDir, "empty-home"),
745+
K_SERVICE: "cloud-run-service",
746+
}),
747+
).toBe(true);
748+
});
749+
750+
it("resolves non-file Vertex ADC through google-auth-library without OAuth refresh fetch", async () => {
751+
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-authlib-"));
752+
vi.stubEnv("GOOGLE_APPLICATION_CREDENTIALS", "");
753+
vi.stubEnv("HOME", path.join(tempDir, "home"));
754+
vi.stubEnv("APPDATA", "");
755+
googleAuthGetAccessTokenMock.mockResolvedValueOnce("ya29.google-auth-token");
756+
const tokenFetchMock = vi.fn();
757+
758+
await expect(resolveGoogleVertexAuthorizedUserHeaders(tokenFetchMock)).resolves.toEqual({
759+
Authorization: "Bearer ya29.google-auth-token",
760+
});
761+
762+
expect(googleAuthMock).toHaveBeenCalledWith({
763+
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
764+
});
765+
expect(googleAuthGetAccessTokenMock).toHaveBeenCalledTimes(1);
766+
expect(tokenFetchMock).not.toHaveBeenCalled();
767+
});
768+
769+
it("uses google-auth-library bearer auth for Google Vertex credential marker requests", async () => {
770+
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-authlib-stream-"));
771+
vi.stubEnv("GOOGLE_APPLICATION_CREDENTIALS", "");
772+
vi.stubEnv("HOME", path.join(tempDir, "home"));
773+
vi.stubEnv("APPDATA", "");
774+
vi.stubEnv("GOOGLE_CLOUD_PROJECT", "vertex-project");
775+
vi.stubEnv("GOOGLE_CLOUD_LOCATION", "us-central1");
776+
googleAuthGetAccessTokenMock.mockResolvedValueOnce("ya29.transport-token");
777+
const tokenFetchMock = vi.fn();
778+
guardedFetchMock.mockResolvedValueOnce(
779+
buildSseResponse([
780+
{
781+
candidates: [{ content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }],
782+
},
783+
]),
784+
);
785+
786+
const streamFn = createGoogleVertexTransportStreamFn();
787+
const stream = await Promise.resolve(
788+
streamFn(
789+
buildGoogleVertexModel(),
790+
{
791+
messages: [{ role: "user", content: "hello", timestamp: 0 }],
792+
} as Parameters<typeof streamFn>[1],
793+
{
794+
apiKey: "gcp-vertex-credentials",
795+
fetch: tokenFetchMock,
796+
} as Parameters<typeof streamFn>[2],
797+
),
798+
);
799+
await stream.result();
800+
801+
expect(tokenFetchMock).not.toHaveBeenCalled();
802+
const guardedCall = requireMockCall(guardedFetchMock, 0, "guarded fetch");
803+
const guardedInit = requireRequestInit(guardedCall, "guarded fetch");
804+
expectHeaders(guardedInit, {
805+
Authorization: "Bearer ya29.transport-token",
806+
"Content-Type": "application/json",
807+
accept: "text/event-stream",
808+
});
809+
expect(new Headers(guardedInit.headers).has("x-goog-api-key")).toBe(false);
810+
});
811+
698812
it("refreshes authorized_user ADC before Google Vertex requests", async () => {
699813
const tempDir = await mkdtemp(path.join(os.tmpdir(), "openclaw-google-vertex-adc-"));
700814
const credentialsPath = path.join(tempDir, "application_default_credentials.json");

‎extensions/google/vertex-adc.ts‎

Lines changed: 146 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,33 @@ type GoogleVertexAuthorizedUserToken = {
1717
refreshToken: string;
1818
};
1919

20+
type GoogleVertexAdcToken = {
21+
token: string;
22+
expiresAtMs: number;
23+
};
24+
2025
const GCP_VERTEX_CREDENTIALS_MARKER = "gcp-vertex-credentials";
2126
const GOOGLE_OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token";
27+
const GOOGLE_VERTEX_OAUTH_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
28+
// Hold tokens slightly less long than reported expiry (Google's recommendation
29+
// is a 60s buffer) so we don't ship a request that's already revoked when it
30+
// leaves the gateway.
31+
const GOOGLE_VERTEX_TOKEN_EXPIRY_BUFFER_MS = 60_000;
2232

2333
let cachedGoogleVertexAuthorizedUserToken: GoogleVertexAuthorizedUserToken | undefined;
34+
let cachedGoogleAuthClient:
35+
| {
36+
promise: Promise<{
37+
getAccessToken: () => Promise<string | null | undefined>;
38+
}>;
39+
}
40+
| undefined;
41+
let cachedGoogleVertexAdcToken: GoogleVertexAdcToken | undefined;
2442

2543
export function resetGoogleVertexAuthorizedUserTokenCacheForTest(): void {
2644
cachedGoogleVertexAuthorizedUserToken = undefined;
45+
cachedGoogleAuthClient = undefined;
46+
cachedGoogleVertexAdcToken = undefined;
2747
}
2848

2949
function normalizeOptionalString(value: unknown): string | undefined {
@@ -85,24 +105,64 @@ async function readGoogleAuthorizedUserCredentials(
85105
};
86106
}
87107

108+
function readGoogleAdcCredentialsTypeSync(credentialsPath: string): string | undefined {
109+
try {
110+
const parsed = JSON.parse(readFileSync(credentialsPath, "utf8")) as unknown;
111+
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
112+
return undefined;
113+
}
114+
const type = (parsed as { type?: unknown }).type;
115+
return typeof type === "string" ? type : undefined;
116+
} catch {
117+
return undefined;
118+
}
119+
}
120+
121+
/**
122+
* Returns true when *any* Application Default Credentials source usable for
123+
* Google Vertex AI is detectable synchronously. We still call the function
124+
* `...AuthorizedUserAdcSync` for backwards compatibility with consumers in
125+
* older OpenClaw revisions; the predicate now also covers:
126+
*
127+
* 1. `authorized_user` credentials file (existing case - `gcloud auth
128+
* application-default login` produces this).
129+
* 2. `external_account` credentials file (Workload Identity Federation).
130+
* 3. `service_account` credentials file (raw GSA key - rarely used in
131+
* OpenClaw, included for completeness).
132+
* 4. GKE Workload Identity (no credentials file; the Google Cloud metadata
133+
* server is the source). Detected heuristically via the presence of
134+
* `KUBERNETES_SERVICE_HOST` (which the kubelet sets in every container).
135+
* The actual metadata-server probe happens at first request time inside
136+
* `google-auth-library`'s `GoogleAuth#getAccessToken()`.
137+
* 5. Cloud Run / GAE / Compute Engine metadata server, detected via
138+
* `K_SERVICE` (Cloud Run), `GAE_SERVICE` (App Engine), or
139+
* `GCE_METADATA_HOST` (Compute Engine override).
140+
*
141+
* The point of the gating in `provider-registration.ts` is to decide whether
142+
* to wire up the Google Vertex transport at all; returning `true` does not
143+
* mean a token is available, it means we have evidence we *could* obtain
144+
* one. The auth call still runs at request time and surfaces clear errors
145+
* when no source is actually usable.
146+
*/
88147
export function hasGoogleVertexAuthorizedUserAdcSync(
89148
env: NodeJS.ProcessEnv = process.env,
90149
): boolean {
91150
const credentialsPath = resolveGoogleApplicationCredentialsPath(env);
92-
if (!credentialsPath) {
93-
return false;
151+
if (credentialsPath) {
152+
const type = readGoogleAdcCredentialsTypeSync(credentialsPath);
153+
if (type === "authorized_user" || type === "external_account" || type === "service_account") {
154+
return true;
155+
}
94156
}
95-
try {
96-
const parsed = JSON.parse(readFileSync(credentialsPath, "utf8")) as unknown;
97-
return (
98-
Boolean(parsed) &&
99-
typeof parsed === "object" &&
100-
!Array.isArray(parsed) &&
101-
(parsed as { type?: unknown }).type === "authorized_user"
102-
);
103-
} catch {
104-
return false;
157+
if (
158+
normalizeOptionalString(env.KUBERNETES_SERVICE_HOST) ||
159+
normalizeOptionalString(env.K_SERVICE) ||
160+
normalizeOptionalString(env.GAE_SERVICE) ||
161+
normalizeOptionalString(env.GCE_METADATA_HOST)
162+
) {
163+
return true;
105164
}
165+
return false;
106166
}
107167

108168
async function refreshGoogleVertexAuthorizedUserAccessToken(params: {
@@ -123,7 +183,7 @@ async function refreshGoogleVertexAuthorizedUserAccessToken(params: {
123183
if (
124184
cached?.credentialsPath === params.credentialsPath &&
125185
cached.refreshToken === refreshToken &&
126-
cached.expiresAtMs - Date.now() > 60_000
186+
cached.expiresAtMs - Date.now() > GOOGLE_VERTEX_TOKEN_EXPIRY_BUFFER_MS
127187
) {
128188
return cached.token;
129189
}
@@ -166,23 +226,83 @@ async function refreshGoogleVertexAuthorizedUserAccessToken(params: {
166226
return token;
167227
}
168228

229+
async function resolveGoogleVertexAccessTokenViaGoogleAuth(): Promise<string> {
230+
// Lazy-import + cache so we don't pay the google-auth-library load cost on
231+
// gateway startup; only when we actually need a non-authorized_user token.
232+
if (!cachedGoogleAuthClient) {
233+
cachedGoogleAuthClient = {
234+
promise: import("google-auth-library").then(({ GoogleAuth }) => {
235+
// GoogleAuth handles every ADC variant we care about for GKE:
236+
// - external_account (Workload Identity Federation: STS exchange)
237+
// - service_account (raw GSA key: JWT-bearer)
238+
// - GKE Workload Identity (metadata server when no credentials file)
239+
// - Compute Engine / Cloud Run / GAE metadata server fallback
240+
// It also caches tokens internally and refreshes before expiry.
241+
return new GoogleAuth({
242+
scopes: [GOOGLE_VERTEX_OAUTH_SCOPE],
243+
});
244+
}),
245+
};
246+
}
247+
const auth = await cachedGoogleAuthClient.promise;
248+
249+
const cached = cachedGoogleVertexAdcToken;
250+
if (cached && cached.expiresAtMs - Date.now() > GOOGLE_VERTEX_TOKEN_EXPIRY_BUFFER_MS) {
251+
return cached.token;
252+
}
253+
254+
const token = await auth.getAccessToken();
255+
const normalized = normalizeOptionalString(token);
256+
if (!normalized) {
257+
throw new Error(
258+
"Google Vertex ADC fallback (google-auth-library) did not return an access token. " +
259+
"Verify the GKE Workload Identity binding (KSA \u2192 GSA), `GOOGLE_APPLICATION_CREDENTIALS`, " +
260+
"or other ADC source is reachable from this pod.",
261+
);
262+
}
263+
// google-auth-library doesn't expose token expiry on the simple
264+
// `getAccessToken()` return type, so we cache for a conservative 5 minutes.
265+
// The library itself already refreshes well before its own internal expiry,
266+
// so this cache is mainly to avoid hot-loop calls into the auth client.
267+
cachedGoogleVertexAdcToken = {
268+
token: normalized,
269+
expiresAtMs: Date.now() + 5 * 60_000,
270+
};
271+
return normalized;
272+
}
273+
274+
/**
275+
* Resolve `Authorization: Bearer ...` headers for Google Vertex calls.
276+
*
277+
* We try the hand-rolled `authorized_user` refresh path first (preserves the
278+
* existing fetchImpl test seam and the OpenClaw upstream behaviour); when the
279+
* configured ADC source is anything other than `authorized_user` (the common
280+
* production cases on GKE: Workload Identity, Workload Identity Federation,
281+
* service-account JSON keys), we hand off to `google-auth-library` which
282+
* understands all of those natively.
283+
*
284+
* Note: the function is still named `...AuthorizedUserHeaders` to avoid a
285+
* symbol rename across the existing patch surface; the docstring above is
286+
* the truth, the name is legacy.
287+
*/
169288
export async function resolveGoogleVertexAuthorizedUserHeaders(
170289
fetchImpl?: typeof fetch,
171290
): Promise<Record<string, string>> {
172291
const credentialsPath = resolveGoogleApplicationCredentialsPath();
173-
if (!credentialsPath) {
174-
throw new Error(
175-
"Google Vertex ADC credentials not found. Set GOOGLE_APPLICATION_CREDENTIALS or run gcloud auth application-default login.",
176-
);
292+
if (credentialsPath) {
293+
const credentials = await readGoogleAuthorizedUserCredentials(credentialsPath);
294+
if (credentials) {
295+
const token = await refreshGoogleVertexAuthorizedUserAccessToken({
296+
credentialsPath,
297+
credentials,
298+
fetchImpl,
299+
});
300+
return { Authorization: `Bearer ${token}` };
301+
}
177302
}
178-
const credentials = await readGoogleAuthorizedUserCredentials(credentialsPath);
179-
if (!credentials) {
180-
throw new Error("Google Vertex ADC fallback requires an authorized_user credentials file.");
181-
}
182-
const token = await refreshGoogleVertexAuthorizedUserAccessToken({
183-
credentialsPath,
184-
credentials,
185-
fetchImpl,
186-
});
303+
// No file-based authorized_user ADC. Fall back to google-auth-library which
304+
// handles GKE Workload Identity (metadata server), Workload Identity
305+
// Federation (external_account), and service-account keys.
306+
const token = await resolveGoogleVertexAccessTokenViaGoogleAuth();
187307
return { Authorization: `Bearer ${token}` };
188308
}

0 commit comments

Comments
 (0)