From 792b3bff3d2f73b8024323776b9450d54c114228 Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Tue, 3 Mar 2026 16:34:42 -0500 Subject: [PATCH] add google credentials --- packages/core/src/agents/agentLoader.ts | 26 ++- .../core/src/agents/auth-provider/factory.ts | 14 +- .../google-credentials-provider.test.ts | 148 ++++++++++++++++ .../google-credentials-provider.ts | 163 ++++++++++++++++++ packages/core/src/agents/registry.ts | 1 + .../core/src/agents/remote-invocation.test.ts | 1 + packages/core/src/agents/remote-invocation.ts | 35 +--- 7 files changed, 349 insertions(+), 39 deletions(-) create mode 100644 packages/core/src/agents/auth-provider/google-credentials-provider.test.ts create mode 100644 packages/core/src/agents/auth-provider/google-credentials-provider.ts diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 6821854ffd..b13b460899 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -44,7 +44,7 @@ interface FrontmatterLocalAgentDefinition * Authentication configuration for remote agents in frontmatter format. */ interface FrontmatterAuthConfig { - type: 'apiKey' | 'http'; + type: 'apiKey' | 'http' | 'google-credentials'; agent_card_requires_auth?: boolean; // API Key key?: string; @@ -55,6 +55,8 @@ interface FrontmatterAuthConfig { username?: string; password?: string; value?: string; + // Google Credentials + scopes?: string[]; } interface FrontmatterRemoteAgentDefinition @@ -147,8 +149,21 @@ const httpAuthSchema = z.object({ value: z.string().min(1).optional(), }); +/** + * Google Credentials auth schema. + */ +const googleCredentialsAuthSchema = z.object({ + ...baseAuthFields, + type: z.literal('google-credentials'), + scopes: z.array(z.string()).optional(), +}); + const authConfigSchema = z - .discriminatedUnion('type', [apiKeyAuthSchema, httpAuthSchema]) + .discriminatedUnion('type', [ + apiKeyAuthSchema, + httpAuthSchema, + googleCredentialsAuthSchema, + ]) .superRefine((data, ctx) => { if (data.type === 'http') { if (data.value) { @@ -348,6 +363,13 @@ function convertFrontmatterAuthToConfig( name: frontmatter.name, }; + case 'google-credentials': + return { + ...base, + type: 'google-credentials', + scopes: frontmatter.scopes, + }; + case 'http': { if (!frontmatter.scheme) { throw new Error( diff --git a/packages/core/src/agents/auth-provider/factory.ts b/packages/core/src/agents/auth-provider/factory.ts index 66b14d0a32..858a6aaefc 100644 --- a/packages/core/src/agents/auth-provider/factory.ts +++ b/packages/core/src/agents/auth-provider/factory.ts @@ -12,12 +12,15 @@ import type { } from './types.js'; import { ApiKeyAuthProvider } from './api-key-provider.js'; import { HttpAuthProvider } from './http-provider.js'; +import { GoogleCredentialsAuthProvider } from './google-credentials-provider.js'; export interface CreateAuthProviderOptions { /** Required for OAuth/OIDC token storage. */ agentName?: string; authConfig?: A2AAuthConfig; agentCard?: AgentCard; + /** Required by some providers (like google-credentials) to determine token audience. */ + targetUrl?: string; } /** @@ -41,9 +44,14 @@ export class A2AAuthProviderFactory { } switch (authConfig.type) { - case 'google-credentials': - // TODO: Implement - throw new Error('google-credentials auth provider not yet implemented'); + case 'google-credentials': { + const provider = new GoogleCredentialsAuthProvider( + authConfig, + options.targetUrl, + ); + await provider.initialize(); + return provider; + } case 'apiKey': { const provider = new ApiKeyAuthProvider(authConfig); diff --git a/packages/core/src/agents/auth-provider/google-credentials-provider.test.ts b/packages/core/src/agents/auth-provider/google-credentials-provider.test.ts new file mode 100644 index 0000000000..2b27e0ff38 --- /dev/null +++ b/packages/core/src/agents/auth-provider/google-credentials-provider.test.ts @@ -0,0 +1,148 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; +import { GoogleCredentialsAuthProvider } from './google-credentials-provider.js'; +import type { GoogleCredentialsAuthConfig } from './types.js'; +import { GoogleAuth } from 'google-auth-library'; +import { OAuthUtils } from '../../mcp/oauth-utils.js'; + +// Mock the external dependencies +vi.mock('google-auth-library', () => ({ + GoogleAuth: vi.fn(), + })); + +describe('GoogleCredentialsAuthProvider', () => { + const mockConfig: GoogleCredentialsAuthConfig = { + type: 'google-credentials', + }; + + let mockGetClient: Mock; + let mockGetAccessToken: Mock; + let mockGetIdTokenClient: Mock; + let mockFetchIdToken: Mock; + + beforeEach(() => { + vi.clearAllMocks(); + + mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'mock-access-token' }); + mockGetClient = vi.fn().mockResolvedValue({ + getAccessToken: mockGetAccessToken, + }); + + mockFetchIdToken = vi.fn().mockResolvedValue('mock-id-token'); + mockGetIdTokenClient = vi.fn().mockResolvedValue({ + idTokenProvider: { + fetchIdToken: mockFetchIdToken, + }, + }); + + (GoogleAuth as Mock).mockImplementation(() => ({ + getClient: mockGetClient, + getIdTokenClient: mockGetIdTokenClient, + })); + }); + + describe('Initialization', () => { + it('throws if no targetUrl is provided', () => { + expect(() => new GoogleCredentialsAuthProvider(mockConfig)).toThrow( + /targetUrl must be provided/, + ); + }); + + it('throws if targetHost is not allowed', () => { + expect( + () => + new GoogleCredentialsAuthProvider(mockConfig, 'https://example.com'), + ).toThrow(/is not an allowed host/); + }); + + it('initializes seamlessly with .googleapis.com', () => { + expect( + () => + new GoogleCredentialsAuthProvider( + mockConfig, + 'https://language.googleapis.com/v1/models', + ), + ).not.toThrow(); + }); + + it('initializes seamlessly with .run.app', () => { + expect( + () => + new GoogleCredentialsAuthProvider( + mockConfig, + 'https://my-cloud-run-service.run.app', + ), + ).not.toThrow(); + }); + + it('initializes seamlessly with .luci.app', () => { + expect( + () => + new GoogleCredentialsAuthProvider( + mockConfig, + 'https://my-service.luci.app', + ), + ).not.toThrow(); + }); + }); + + describe('Token Fetching', () => { + it('fetches an access token for googleapis.com endpoint', async () => { + const provider = new GoogleCredentialsAuthProvider( + mockConfig, + 'https://language.googleapis.com', + ); + const headers = await provider.headers(); + + expect(headers).toEqual({ Authorization: 'Bearer mock-access-token' }); + expect(mockGetClient).toHaveBeenCalled(); + expect(mockGetAccessToken).toHaveBeenCalled(); + expect(mockGetIdTokenClient).not.toHaveBeenCalled(); + }); + + it('fetches an identity token for run.app endpoint', async () => { + // Mock OAuthUtils.parseTokenExpiry to avoid Base64 decoding issues in tests + vi.spyOn(OAuthUtils, 'parseTokenExpiry').mockReturnValue( + Date.now() + 1000000, + ); + + const provider = new GoogleCredentialsAuthProvider( + mockConfig, + 'https://my-service.run.app/some-path', + ); + const headers = await provider.headers(); + + expect(headers).toEqual({ Authorization: 'Bearer mock-id-token' }); + expect(mockGetIdTokenClient).toHaveBeenCalledWith('my-service.run.app'); + expect(mockFetchIdToken).toHaveBeenCalledWith('my-service.run.app'); + expect(mockGetClient).not.toHaveBeenCalled(); + }); + + it('re-fetches token on auth failure (shouldRetryWithHeaders)', async () => { + vi.spyOn(OAuthUtils, 'parseTokenExpiry').mockReturnValue( + Date.now() + 1000000, + ); + const provider = new GoogleCredentialsAuthProvider( + mockConfig, + 'https://language.googleapis.com', + ); + + const req = {} as RequestInit; + const res = { status: 401 } as Response; + + const retryHeaders = await provider.shouldRetryWithHeaders(req, res); + + expect(retryHeaders).toEqual({ + Authorization: 'Bearer mock-access-token', + }); + expect(mockGetAccessToken).toHaveBeenCalledTimes(1); // the retry fetched it + }); + }); +}); diff --git a/packages/core/src/agents/auth-provider/google-credentials-provider.ts b/packages/core/src/agents/auth-provider/google-credentials-provider.ts new file mode 100644 index 0000000000..ac8211d2f0 --- /dev/null +++ b/packages/core/src/agents/auth-provider/google-credentials-provider.ts @@ -0,0 +1,163 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { HttpHeaders } from '@a2a-js/sdk/client'; +import { BaseA2AAuthProvider } from './base-provider.js'; +import type { GoogleCredentialsAuthConfig } from './types.js'; +import { GoogleAuth } from 'google-auth-library'; +import { debugLogger } from '../../utils/debugLogger.js'; +import { OAuthUtils, FIVE_MIN_BUFFER_MS } from '../../mcp/oauth-utils.js'; + +const CLOUD_RUN_HOST_REGEX = /^(.*\.)?run\.app$/; +const CLOUD_LUCI_HOST_REGEX = /^(.*\.)?luci\.app$/; +const ALLOWED_HOSTS = [ + /^.+\.googleapis\.com$/, + CLOUD_LUCI_HOST_REGEX, + CLOUD_RUN_HOST_REGEX, +]; + +/** + * Authentication provider for Google ADC (Application Default Credentials). + * Automatically decides whether to use identity tokens or access tokens + * based on the target endpoint URL. + */ +export class GoogleCredentialsAuthProvider extends BaseA2AAuthProvider { + readonly type = 'google-credentials' as const; + + private readonly auth: GoogleAuth; + private readonly useIdToken: boolean = false; + private readonly audience?: string; + private cachedToken?: string; + private tokenExpiryTime?: number; + + constructor( + private readonly config: GoogleCredentialsAuthConfig, + targetUrl?: string, + ) { + super(); + + if (!targetUrl) { + throw new Error( + 'targetUrl must be provided to GoogleCredentialsAuthProvider to determine token audience.', + ); + } + + const hostname = new URL(targetUrl).hostname; + const isRunAppHost = CLOUD_RUN_HOST_REGEX.test(hostname); + const isLuciAppHost = CLOUD_LUCI_HOST_REGEX.test(hostname); + + if (isRunAppHost || isLuciAppHost) { + this.useIdToken = true; + } + this.audience = hostname; + + if ( + !this.useIdToken && + !ALLOWED_HOSTS.some((pattern) => pattern.test(hostname)) + ) { + throw new Error( + `Host "${hostname}" is not an allowed host for Google Credential provider.`, + ); + } + + // A2A spec requires scopes if configured, otherwise use default cloud-platform + const scopes = + this.config.scopes && this.config.scopes.length > 0 + ? this.config.scopes + : ['https://www.googleapis.com/auth/cloud-platform']; + + this.auth = new GoogleAuth({ + scopes, + }); + } + + override async initialize(): Promise { + // We can pre-fetch or validate if necessary here, + // but deferred fetching is usually better for auth tokens. + } + + async headers(): Promise { + // Check cache + if ( + this.cachedToken && + this.tokenExpiryTime && + Date.now() < this.tokenExpiryTime - FIVE_MIN_BUFFER_MS + ) { + return { Authorization: `Bearer ${this.cachedToken}` }; + } + + // Clear expired cache + this.cachedToken = undefined; + this.tokenExpiryTime = undefined; + + if (this.useIdToken) { + try { + const idClient = await this.auth.getIdTokenClient(this.audience!); + const idToken = await idClient.idTokenProvider.fetchIdToken( + this.audience!, + ); + + const expiryTime = OAuthUtils.parseTokenExpiry(idToken); + if (expiryTime) { + this.tokenExpiryTime = expiryTime; + this.cachedToken = idToken; + } + + return { Authorization: `Bearer ${idToken}` }; + } catch (e) { + const errorMessage = `Failed to get ADC ID token: ${ + e instanceof Error ? e.message : String(e) + }`; + debugLogger.error(errorMessage, e); + throw new Error(errorMessage); + } + } + + // Otherwise, access token + try { + const client = await this.auth.getClient(); + const token = await client.getAccessToken(); + + if (token.token) { + // We do not cache the access token ourselves since GoogleAuth Client does it under the hood. + // Or if we wanted to, we could, but let's rely on getAccessToken()'s cache for now. + return { Authorization: `Bearer ${token.token}` }; + } + throw new Error('Failed to retrieve ADC access token.'); + } catch (e) { + const errorMessage = `Failed to get ADC access token: ${ + e instanceof Error ? e.message : String(e) + }`; + debugLogger.error(errorMessage, e); + throw new Error(errorMessage); + } + } + + override async shouldRetryWithHeaders( + _req: RequestInit, + res: Response, + ): Promise { + if (res.status !== 401 && res.status !== 403) { + this.authRetryCount = 0; + return undefined; + } + + if (this.authRetryCount >= BaseA2AAuthProvider.MAX_AUTH_RETRIES) { + return undefined; + } + this.authRetryCount++; + + debugLogger.debug( + '[GoogleCredentialsAuthProvider] Re-fetching token after auth failure', + ); + + // Clear cache to force a re-fetch + this.cachedToken = undefined; + this.tokenExpiryTime = undefined; + + return this.headers(); + } +} diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index bf7e669150..4f7fda2479 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -416,6 +416,7 @@ export class AgentRegistry { const provider = await A2AAuthProviderFactory.create({ authConfig: definition.auth, agentName: definition.name, + targetUrl: definition.agentCardUrl, }); if (!provider) { throw new Error( diff --git a/packages/core/src/agents/remote-invocation.test.ts b/packages/core/src/agents/remote-invocation.test.ts index 02c655ec27..ccf43a2165 100644 --- a/packages/core/src/agents/remote-invocation.test.ts +++ b/packages/core/src/agents/remote-invocation.test.ts @@ -195,6 +195,7 @@ describe('RemoteAgentInvocation', () => { expect(A2AAuthProviderFactory.create).toHaveBeenCalledWith({ authConfig: mockAuth, agentName: 'test-agent', + targetUrl: 'http://test-agent/card', }); expect(mockClientManager.loadAgent).toHaveBeenCalledWith( 'test-agent', diff --git a/packages/core/src/agents/remote-invocation.ts b/packages/core/src/agents/remote-invocation.ts index dad7f8167d..6888d97960 100644 --- a/packages/core/src/agents/remote-invocation.ts +++ b/packages/core/src/agents/remote-invocation.ts @@ -19,46 +19,12 @@ import type { import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { A2AClientManager } from './a2a-client-manager.js'; import { extractIdsFromResponse, A2AResultReassembler } from './a2aUtils.js'; -import { GoogleAuth } from 'google-auth-library'; import type { AuthenticationHandler } from '@a2a-js/sdk/client'; import { debugLogger } from '../utils/debugLogger.js'; import type { AnsiOutput } from '../utils/terminalSerializer.js'; import type { SendMessageResult } from './a2a-client-manager.js'; import { A2AAuthProviderFactory } from './auth-provider/factory.js'; -/** - * Authentication handler implementation using Google Application Default Credentials (ADC). - */ -export class ADCHandler implements AuthenticationHandler { - private auth = new GoogleAuth({ - scopes: ['https://www.googleapis.com/auth/cloud-platform'], - }); - - async headers(): Promise> { - try { - const client = await this.auth.getClient(); - const token = await client.getAccessToken(); - if (token.token) { - return { Authorization: `Bearer ${token.token}` }; - } - throw new Error('Failed to retrieve ADC access token.'); - } catch (e) { - const errorMessage = `Failed to get ADC token: ${ - e instanceof Error ? e.message : String(e) - }`; - debugLogger.log('ERROR', errorMessage); - throw new Error(errorMessage); - } - } - - async shouldRetryWithHeaders( - _response: unknown, - ): Promise | undefined> { - // For ADC, we usually just re-fetch the token if needed. - return this.headers(); - } -} - /** * A tool invocation that proxies to a remote A2A agent. * @@ -117,6 +83,7 @@ export class RemoteAgentInvocation extends BaseToolInvocation< const provider = await A2AAuthProviderFactory.create({ authConfig: this.definition.auth, agentName: this.definition.name, + targetUrl: this.definition.agentCardUrl, }); if (!provider) { throw new Error(