diff --git a/packages/ai/src/bootstrap/main.ts b/packages/ai/src/bootstrap/main.ts index b65e833a..c7488adc 100644 --- a/packages/ai/src/bootstrap/main.ts +++ b/packages/ai/src/bootstrap/main.ts @@ -10,6 +10,7 @@ export interface AIGatewayConfig { env: Record siteID: string | undefined siteURL: string | undefined + existingToken?: string } export interface AIProviderEnvVar { @@ -99,9 +100,11 @@ export const fetchAIProviders = async ({ api }: { api: NetlifyAPI }): Promise => { try { if (!api.accessToken) { @@ -111,12 +114,18 @@ export const fetchAIGatewayToken = async ({ // TODO: update once available in openApi const url = `${api.scheme}://${api.host}/api/v1/sites/${siteId}/ai-gateway/token` + const headers: Record = { + Authorization: `Bearer ${api.accessToken}`, + 'Content-Type': 'application/json', + } + + if (priorAuthToken) { + headers['X-Prior-Authorization'] = priorAuthToken + } + const response = await fetch(url, { method: 'GET', - headers: { - Authorization: `Bearer ${api.accessToken}`, - 'Content-Type': 'application/json', - }, + headers, }) if (!response.ok) { @@ -144,12 +153,26 @@ export const fetchAIGatewayToken = async ({ } } -export const setupAIGateway = async (config: AIGatewayConfig): Promise => { - const { api, env, siteID, siteURL } = config +export const setupAIGateway = async (config: AIGatewayConfig): Promise<{ token: string; url: string } | null> => { + const { api, env, siteID, siteURL, existingToken } = config if (siteID && siteID !== 'unlinked' && siteURL) { + let priorAuthToken: string | undefined + + // If existingToken is explicitly provided (even if empty string), use it + if (existingToken !== undefined) { + priorAuthToken = existingToken || undefined + } else { + // If no existingToken provided, extract existing AI_GATEWAY from process.env to check for prior auth token + const existingAIGateway = parseAIGatewayContext(process.env.AI_GATEWAY) + // If there's an existing AI Gateway context with the same URL, use its token as prior auth + if (existingAIGateway && existingAIGateway.url === `${siteURL}/.netlify/ai`) { + priorAuthToken = existingAIGateway.token + } + } + const [aiGatewayToken, envVars] = await Promise.all([ - fetchAIGatewayToken({ api, siteId: siteID }), + fetchAIGatewayToken({ api, siteId: siteID, priorAuthToken }), fetchAIProviders({ api }), ]) @@ -161,8 +184,15 @@ export const setupAIGateway = async (config: AIGatewayConfig): Promise => }) const base64Context = Buffer.from(aiGatewayContext).toString('base64') env.AI_GATEWAY = { sources: ['internal'], value: base64Context } + + return { + token: aiGatewayToken.token, + url: `${siteURL}/.netlify/ai`, + } } } + + return null } export const parseAIGatewayContext = (aiGatewayValue?: string): AIGatewayTokenResponse | undefined => { diff --git a/packages/ai/src/main.test.ts b/packages/ai/src/main.test.ts index ac613e4f..69ae9cb9 100644 --- a/packages/ai/src/main.test.ts +++ b/packages/ai/src/main.test.ts @@ -1,4 +1,4 @@ -import { describe, test, expect, vi, beforeEach } from 'vitest' +import { describe, test, expect, vi, beforeEach, afterEach } from 'vitest' import { fetchAIGatewayToken, setupAIGateway, parseAIGatewayContext, fetchAIProviders } from './bootstrap/main.js' import type { NetlifyAPI } from '@netlify/api' @@ -42,6 +42,34 @@ describe('fetchAIGatewayToken', () => { }) }) + test('successfully fetches AI Gateway token with prior authorization', async () => { + const mockResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai/', + } + + mockFetch.mockResolvedValue({ + ok: true, + json: () => Promise.resolve(mockResponse), + }) + + const result = await fetchAIGatewayToken({ + api: mockApi, + siteId: 'test-site-id', + priorAuthToken: 'prior-token', + }) + + expect(result).toEqual(mockResponse) + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site-id/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + 'X-Prior-Authorization': 'prior-token', + }, + }) + }) + test('returns null when no access token is provided', async () => { const apiWithoutToken: NetlifyAPI = { scheme: mockApi.scheme, @@ -231,8 +259,17 @@ describe('setupAIGateway', () => { accessToken: 'test-token', } as NetlifyAPI + const originalProcessEnv = process.env + beforeEach(() => { vi.clearAllMocks() + // Reset process.env to original state + process.env = { ...originalProcessEnv } + }) + + afterEach(() => { + // Restore original process.env + process.env = originalProcessEnv }) test('sets up AI Gateway when conditions are met', async () => { @@ -310,6 +347,284 @@ describe('setupAIGateway', () => { expect(env).not.toHaveProperty('AI_GATEWAY') }) + + test('uses prior authorization token from existing AI_GATEWAY in process.env', async () => { + const existingContext = { + token: 'existing-token', + url: 'https://example.com/.netlify/ai', + envVars: [], + } + const existingBase64 = Buffer.from(JSON.stringify(existingContext)).toString('base64') + process.env.AI_GATEWAY = existingBase64 + + const mockTokenResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai', + } + + const mockProvidersResponse = { + providers: { + openai: { + token_env_var: 'OPENAI_API_KEY', + url_env_var: 'OPENAI_BASE_URL', + models: ['gpt-4'], + }, + }, + } + + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockTokenResponse), + }) + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockProvidersResponse), + }) + + const env = {} + const config = { + api: mockApi, + env, + siteID: 'test-site', + siteURL: 'https://example.com', + } + + await setupAIGateway(config) + + // Verify that the fetchAIGatewayToken was called with the prior auth token + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + 'X-Prior-Authorization': 'existing-token', + }, + }) + + expect(env).toHaveProperty('AI_GATEWAY') + }) + + test('does not use prior authorization when existing AI_GATEWAY has different URL', async () => { + const existingContext = { + token: 'existing-token', + url: 'https://different-site.com/.netlify/ai', + envVars: [], + } + const existingBase64 = Buffer.from(JSON.stringify(existingContext)).toString('base64') + process.env.AI_GATEWAY = existingBase64 + + const mockTokenResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai', + } + + const mockProvidersResponse = { + providers: { + openai: { + token_env_var: 'OPENAI_API_KEY', + url_env_var: 'OPENAI_BASE_URL', + models: ['gpt-4'], + }, + }, + } + + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockTokenResponse), + }) + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockProvidersResponse), + }) + + const env = {} + const config = { + api: mockApi, + env, + siteID: 'test-site', + siteURL: 'https://example.com', + } + + await setupAIGateway(config) + + // Verify that the fetchAIGatewayToken was called without prior auth token + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + }, + }) + + expect(env).toHaveProperty('AI_GATEWAY') + }) + + test('handles invalid AI_GATEWAY in process.env gracefully', async () => { + process.env.AI_GATEWAY = 'invalid-base64-data' + + const mockTokenResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai', + } + + const mockProvidersResponse = { + providers: { + openai: { + token_env_var: 'OPENAI_API_KEY', + url_env_var: 'OPENAI_BASE_URL', + models: ['gpt-4'], + }, + }, + } + + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockTokenResponse), + }) + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockProvidersResponse), + }) + + const env = {} + const config = { + api: mockApi, + env, + siteID: 'test-site', + siteURL: 'https://example.com', + } + + await setupAIGateway(config) + + // Verify that the fetchAIGatewayToken was called without prior auth token + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + }, + }) + + expect(env).toHaveProperty('AI_GATEWAY') + }) + + test('uses existingToken parameter when provided, ignoring process.env', async () => { + const existingContext = { + token: 'existing-token-from-env', + url: 'https://example.com/.netlify/ai', + envVars: [], + } + const existingBase64 = Buffer.from(JSON.stringify(existingContext)).toString('base64') + process.env.AI_GATEWAY = existingBase64 + + const mockTokenResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai', + } + + const mockProvidersResponse = { + providers: { + openai: { + token_env_var: 'OPENAI_API_KEY', + url_env_var: 'OPENAI_BASE_URL', + models: ['gpt-4'], + }, + }, + } + + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockTokenResponse), + }) + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockProvidersResponse), + }) + + const env = {} + const config = { + api: mockApi, + env, + siteID: 'test-site', + siteURL: 'https://example.com', + existingToken: 'explicit-prior-token', + } + + await setupAIGateway(config) + + // Verify that the fetchAIGatewayToken was called with the explicit existingToken, not the one from env + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + 'X-Prior-Authorization': 'explicit-prior-token', + }, + }) + + expect(env).toHaveProperty('AI_GATEWAY') + }) + + test('uses empty existingToken parameter over process.env when explicitly set to empty string', async () => { + const existingContext = { + token: 'existing-token-from-env', + url: 'https://example.com/.netlify/ai', + envVars: [], + } + const existingBase64 = Buffer.from(JSON.stringify(existingContext)).toString('base64') + process.env.AI_GATEWAY = existingBase64 + + const mockTokenResponse = { + token: 'new-ai-gateway-token', + url: 'https://ai-gateway.com/.netlify/ai', + } + + const mockProvidersResponse = { + providers: { + openai: { + token_env_var: 'OPENAI_API_KEY', + url_env_var: 'OPENAI_BASE_URL', + models: ['gpt-4'], + }, + }, + } + + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockTokenResponse), + }) + .mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve(mockProvidersResponse), + }) + + const env = {} + const config = { + api: mockApi, + env, + siteID: 'test-site', + siteURL: 'https://example.com', + existingToken: '', + } + + await setupAIGateway(config) + + // Verify that the fetchAIGatewayToken was called without prior auth token when existingToken is empty string + expect(mockFetch).toHaveBeenCalledWith('https://api.netlify.com/api/v1/sites/test-site/ai-gateway/token', { + method: 'GET', + headers: { + Authorization: 'Bearer test-token', + 'Content-Type': 'application/json', + }, + }) + + expect(env).toHaveProperty('AI_GATEWAY') + }) }) describe('parseAIGatewayContext', () => {