Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,27 +1,32 @@
import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';

let authInstance: GoogleAuth | null = null;
let authOptions: GoogleAuthOptions | null = null;
/**
* Creates a auth token generator function for the given options.
* This avoids reference-equality cache invalidation issues.
*/
export function createAuthTokenGenerator(options?: GoogleAuthOptions) {
const auth = new GoogleAuth({
scopes: ['https://www.googleapis.com/auth/cloud-platform'],
...options,
});

function getAuth(options: GoogleAuthOptions) {
if (!authInstance || options !== authOptions) {
authInstance = new GoogleAuth({
scopes: ['https://www.googleapis.com/auth/cloud-platform'],
...options,
});
authOptions = options;
}
return authInstance;
return async function generateAuthToken() {
const client = await auth.getClient();
const token = await client.getAccessToken();
return token?.token ?? null;
};
}

/**
* @deprecated Use createAuthTokenGenerator instead.
* This function is kept for backward compatibility.
*/
export async function generateAuthToken(options?: GoogleAuthOptions) {
const auth = getAuth(options || {});
const auth = new GoogleAuth({
scopes: ['https://www.googleapis.com/auth/cloud-platform'],
...options,
});
const client = await auth.getClient();
const token = await client.getAccessToken();
return token?.token || null;
}

// For testing purposes only
export function _resetAuthInstance() {
authInstance = null;
}
}
11 changes: 6 additions & 5 deletions packages/google-vertex/src/google-vertex-provider-node.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { loadOptionalSetting, resolve } from '@ai-sdk/provider-utils';
import { GoogleAuthOptions } from 'google-auth-library';
import { generateAuthToken } from './google-vertex-auth-google-auth-library';
import { createAuthTokenGenerator } from './google-vertex-auth-google-auth-library';
import {
createVertex as createVertexOriginal,
GoogleVertexProvider,
Expand Down Expand Up @@ -28,16 +28,17 @@ export function createVertex(
environmentVariableName: 'GOOGLE_VERTEX_API_KEY',
});

// Create a per-instance auth token generator to avoid reference-equality issues
const generateAuthToken = createAuthTokenGenerator(options.googleAuthOptions);

if (apiKey) {
return createVertexOriginal(options);
}

return createVertexOriginal({
...options,
headers: async () => ({
Authorization: `Bearer ${await generateAuthToken(
options.googleAuthOptions,
)}`,
Authorization: `Bearer ${await generateAuthToken()}`,
...(await resolve(options.headers)),
}),
});
Expand All @@ -46,4 +47,4 @@ export function createVertex(
/**
* Default Google Vertex AI provider instance.
*/
export const vertex = createVertex();
export const vertex = createVertex();