From 8610d384aefa2ced1006cd8ea6789f9243120fc9 Mon Sep 17 00:00:00 2001 From: Sean Roberts Date: Sat, 4 Oct 2025 10:49:42 -0400 Subject: [PATCH 1/3] feat: support prior ips for AI gateway validations --- packages/ai/src/bootstrap/main.ts | 27 +++- packages/ai/src/main.test.ts | 211 +++++++++++++++++++++++++++++- 2 files changed, 232 insertions(+), 6 deletions(-) diff --git a/packages/ai/src/bootstrap/main.ts b/packages/ai/src/bootstrap/main.ts index b65e833a..2e6b103e 100644 --- a/packages/ai/src/bootstrap/main.ts +++ b/packages/ai/src/bootstrap/main.ts @@ -99,9 +99,11 @@ export const fetchAIProviders = async ({ api }: { api: NetlifyAPI }): Promise => { try { if (!api.accessToken) { @@ -111,12 +113,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) { @@ -148,8 +156,17 @@ export const setupAIGateway = async (config: AIGatewayConfig): Promise => const { api, env, siteID, siteURL } = config if (siteID && siteID !== 'unlinked' && siteURL) { + // Extract existing AI_GATEWAY from process.env to check for prior auth token + const existingAIGateway = parseAIGatewayContext(process.env.AI_GATEWAY) + let priorAuthToken: string | undefined + + // 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 }), ]) diff --git a/packages/ai/src/main.test.ts b/packages/ai/src/main.test.ts index ac613e4f..90643347 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,178 @@ 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') + }) }) describe('parseAIGatewayContext', () => { From 104dd648d05ac39580fc2215bdcfd43ab1d00f2a Mon Sep 17 00:00:00 2001 From: Sean Roberts Date: Sat, 4 Oct 2025 10:59:18 -0400 Subject: [PATCH 2/3] feat: return the values from ai setup and allow prior tokens to be sent back in --- packages/ai/src/bootstrap/main.ts | 27 +++++-- packages/ai/src/main.test.ts | 121 ++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/packages/ai/src/bootstrap/main.ts b/packages/ai/src/bootstrap/main.ts index 2e6b103e..c62f4b8b 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 { @@ -152,17 +153,22 @@ 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) { - // Extract existing AI_GATEWAY from process.env to check for prior auth token - const existingAIGateway = parseAIGatewayContext(process.env.AI_GATEWAY) let priorAuthToken: string | undefined - // 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 + // 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([ @@ -178,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 90643347..a1b75a67 100644 --- a/packages/ai/src/main.test.ts +++ b/packages/ai/src/main.test.ts @@ -519,6 +519,127 @@ describe('setupAIGateway', () => { 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', () => { From 2b4d14327f0f381d5a6d35c3f2ecb26ef715cedf Mon Sep 17 00:00:00 2001 From: Sean Roberts Date: Sat, 4 Oct 2025 10:59:59 -0400 Subject: [PATCH 3/3] fix: format --- packages/ai/src/bootstrap/main.ts | 4 +- packages/ai/src/main.test.ts | 89 +++++++++++++------------------ 2 files changed, 39 insertions(+), 54 deletions(-) diff --git a/packages/ai/src/bootstrap/main.ts b/packages/ai/src/bootstrap/main.ts index c62f4b8b..c7488adc 100644 --- a/packages/ai/src/bootstrap/main.ts +++ b/packages/ai/src/bootstrap/main.ts @@ -184,14 +184,14 @@ export const setupAIGateway = async (config: AIGatewayConfig): Promise<{ token: }) const base64Context = Buffer.from(aiGatewayContext).toString('base64') env.AI_GATEWAY = { sources: ['internal'], value: base64Context } - + return { token: aiGatewayToken.token, url: `${siteURL}/.netlify/ai`, } } } - + return null } diff --git a/packages/ai/src/main.test.ts b/packages/ai/src/main.test.ts index a1b75a67..69ae9cb9 100644 --- a/packages/ai/src/main.test.ts +++ b/packages/ai/src/main.test.ts @@ -393,17 +393,14 @@ describe('setupAIGateway', () => { 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(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') }) @@ -453,16 +450,13 @@ describe('setupAIGateway', () => { 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(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') }) @@ -506,16 +500,13 @@ describe('setupAIGateway', () => { 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(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') }) @@ -566,17 +557,14 @@ describe('setupAIGateway', () => { 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(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') }) @@ -627,16 +615,13 @@ describe('setupAIGateway', () => { 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(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') })