Allow vertex to work on Vercel-Edge

[cosbgn]

Feature Description

The vertex model doesn't work on vercel edge. This is cause by the auth-library not being compatible with the edge. This could be easily replaced by a vanilla js script:

import { webcrypto } from 'crypto'
const tokenUrl = 'https://oauth2.googleapis.com/token';
const base64url = (source) => {
    let encodedSource = Buffer.from(source).toString('base64');
    encodedSource = encodedSource.replace(/=+$/, '');
    encodedSource = encodedSource.replace(/\+/g, '-');
    encodedSource = encodedSource.replace(/\//g, '_');
    return encodedSource;
  }

  const createJwtSegment = (segment) => base64url(JSON.stringify(segment));

  const importPrivateKey = async (pemKey) => {
    const pemHeader = "-----BEGIN PRIVATE KEY-----";
    const pemFooter = "-----END PRIVATE KEY-----";
    const pemContents = pemKey.substring(pemHeader.length, pemKey.length - pemFooter.length);
    const binaryDerString = Buffer.from(pemContents, 'base64').toString('binary');
    const binaryDer = new Uint8Array(binaryDerString.length);

    for (let i = 0; i < binaryDerString.length; i++) {
      binaryDer[i] = binaryDerString.charCodeAt(i);
    }

    return webcrypto.subtle.importKey("pkcs8", binaryDer.buffer, { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" }, true, ["sign"] );
  }

  const signJwt = async (header, payload, privateKey) => {
    const data = `${header}.${payload}`;
    const encoder = new TextEncoder();
    const signature = await webcrypto.subtle.sign("RSASSA-PKCS1-v1_5", privateKey, encoder.encode(data) );
    return `${data}.${base64url(Buffer.from(signature))}`;
  };

  const buildJwt = async ({ client_email, private_key, private_key_id, scopes }) => {
    const iat = Math.floor(Date.now() / 1000);
    const exp = iat + 3600;

    const payload = { iss: client_email, sub: client_email, aud: tokenUrl, iat: iat, exp: exp, scope: scopes.join(" ") }
    const jwtHeaders = { kid: private_key_id, alg: 'RS256', typ: 'JWT' };

    const encodedHeader = createJwtSegment(jwtHeaders);
    const encodedPayload = createJwtSegment(payload);

    // Import the private key
    const privateKey = await importPrivateKey(private_key);

    // Sign the JWT
    const signedJwt = await signJwt(encodedHeader, encodedPayload, privateKey);

    return signedJwt;
}

// Function to get access token
const generateToken = async(credentials) => {
    // 1) Build and sign JWT
    const signedJwt = await buildJwt(credentials)
    // 2) Get access token
    const response = await fetch(tokenUrl, {
      method: 'POST',
      headers: {  'Content-Type': 'application/x-www-form-urlencoded' },
      body: new URLSearchParams({
        'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
        'assertion': signedJwt
      })
    })

    if (!response.ok) {
      throw new Error(`Token request failed: ${response.statusText}`);
    }
    return await response.json()
}

It can easily be used like this:

    const { expires_in, access_token } = await generateToken({
            client_email:process.env.GOOGLE_CLOUD_CLIENT_EMAIL, 
            private_key: process.env.GOOGLE_CLOUD_PRIVATE_KEY?.replace(/\\n/gm, "\n"), // Fix issue with new line
            private_key_id: process.env.GOOGLE_CLOUD_PRIVATE_KEY_ID, 
            scopes: ["https://www.googleapis.com/auth/cloud-platform"]
        })

Use Case

Would be great to decouple vertex and auto so that we can use our custom auth script like the one above or simply add the script above directly in the ai-sdk model so that it would work on vercel edge.

Additional context

No response

[cosbgn]

Also vertex has a very simply rest API, so most likely the whole ai-sdk/vertex could only depend on built in fetch and work perfectly everywhere.

[benaduggan]

I wonder if we swap to using dvsekhvalnov/jose-jwt to get it working. IIRC Jose works on the edge

[rszemplinski]

I wonder if we swap to using dvsekhvalnov/jose-jwt to get it working. IIRC Jose works on the edge

Yep, it does according to their README and this PR