feat: support ai gateway for local edge functions (#7637)

* feat: support ai gateway for local edge functions

* update to published dep

* package lock

Author: kathmbeck

package-lock.json

@@ -65,9 +65,8 @@
     "@netlify/build-info": "10.0.8",
     "@netlify/config": "24.0.4",
     "@netlify/dev-utils": "4.2.0",
-    "@netlify/edge-bundler": "14.5.5",
+    "@netlify/edge-bundler": "^14.5.6",
     "@netlify/edge-functions": "2.18.1",
-    "@netlify/edge-functions-bootstrap": "2.14.0",
     "@netlify/headers-parser": "9.0.2",
     "@netlify/local-functions-proxy": "2.0.3",
     "@netlify/redirect-parser": "15.0.3",
@@ -159,6 +158,7 @@
     "@bugsnag/js": "8.4.0",
     "@eslint/compat": "1.3.2",
     "@eslint/js": "9.24.0",
+    "@netlify/edge-functions-bootstrap": "^2.17.1",
     "@netlify/functions": "3.0.4",
     "@netlify/types": "2.0.3",
     "@sindresorhus/slugify": "2.2.1",
@@ -170,6 +170,7 @@
     "@types/envinfo": "7.8.4",
     "@types/eslint-config-prettier": "6.11.3",
     "@types/etag": "1.8.4",
+    "@types/express": "^4.17.23",
     "@types/folder-walker": "3.2.4",
     "@types/gitconfiglocal": "2.0.3",
     "@types/inquirer": "9.0.9",

package.json

@@ -257,6 +257,7 @@ export const dev = async (options: OptionValues, command: BaseCommand) => {
 
   await startProxyServer({
     addonsUrls,
+    aiGatewayContext,
     api,
     blobsContext,
     command,

src/commands/dev/dev.ts

@@ -1,6 +1,7 @@
 import { Buffer } from 'buffer'
 
 export const headers = {
+  AIGateway: 'x-nf-ai-gateway',
   BlobsInfo: 'x-nf-blobs-info',
   DeployID: 'x-nf-deploy-id',
   DeployContext: 'x-nf-deploy-context',

src/lib/edge-functions/headers.ts

@@ -4,6 +4,7 @@ import type { IncomingMessage, ClientRequest } from 'http'
 import { join, resolve } from 'path'
 
 import * as bundler from '@netlify/edge-bundler'
+import type { AIGatewayContext } from '@netlify/ai/bootstrap'
 import getAvailablePort from 'get-port'
 
 import type BaseCommand from '../../commands/base-command.js'
@@ -75,6 +76,7 @@ const createAccountInfoHeader = ({ id }: { id: string }) => {
 
 export const initializeProxy = async ({
   accountId,
+  aiGatewayContext,
   blobsContext,
   command,
   config,
@@ -95,6 +97,7 @@ export const initializeProxy = async ({
   state,
 }: {
   accountId: string
+  aiGatewayContext?: AIGatewayContext | null
   blobsContext: BlobsContextWithEdgeAccess
   command: BaseCommand
   config: NormalizedCachedConfigConfig
@@ -124,6 +127,7 @@ export const initializeProxy = async ({
   // the network if needed. We don't want to wait for that to be completed, or
   // the command will be left hanging.
   const server = prepareServer({
+    aiGatewayContext,
     command,
     config,
     configPath,
@@ -162,6 +166,10 @@ export const initializeProxy = async ({
       ).toString('base64')
     }
 
+    if (aiGatewayContext) {
+      req.headers[headers.AIGateway] = Buffer.from(JSON.stringify(aiGatewayContext)).toString('base64')
+    }
+
     await registry.initialize()
 
     const url = new URL(req.url!, `http://${LOCAL_HOST}:${mainPort}`)
@@ -194,6 +202,7 @@ export const isEdgeFunctionsRequest = (req: IncomingMessage): req is ExtendedInc
   Object.hasOwn(req, headersSymbol)
 
 const prepareServer = async ({
+  aiGatewayContext,
   command,
   config,
   configPath,
@@ -207,6 +216,7 @@ const prepareServer = async ({
   projectDir,
   repositoryRoot,
 }: {
+  aiGatewayContext?: AIGatewayContext | null
   command: BaseCommand
   config: NormalizedCachedConfigConfig
   configPath: string
@@ -245,8 +255,9 @@ const prepareServer = async ({
       servePath,
     })
     const registry = new EdgeFunctionsRegistry({
-      command,
+      aiGatewayContext,
       bundler,
+      command,
       config,
       configPath,
       debug,

src/lib/edge-functions/proxy.ts

@@ -3,6 +3,7 @@ import { join } from 'path'
 import { fileURLToPath } from 'url'
 
 import type { Declaration, EdgeFunction, FunctionConfig, Manifest, ModuleGraph } from '@netlify/edge-bundler'
+import type { AIGatewayContext } from '@netlify/ai/bootstrap'
 import { watchDebounced } from '@netlify/dev-utils'
 
 import BaseCommand from '../../commands/base-command.js'
@@ -31,19 +32,20 @@ type RunIsolate = Awaited<ReturnType<typeof import('@netlify/edge-bundler').serv
 type ModuleJson = ModuleGraph['modules'][number]
 
 interface EdgeFunctionsRegistryOptions {
-  command: BaseCommand
+  aiGatewayContext?: AIGatewayContext | null
   bundler: typeof import('@netlify/edge-bundler')
+  command: BaseCommand
   config: NormalizedCachedConfigConfig
   configPath: string
   debug: boolean
   directories: string[]
   env: Record<string, { sources: string[]; value: string }>
   featureFlags: FeatureFlags
   getUpdatedConfig: () => Promise<NormalizedCachedConfigConfig>
+  importMapFromTOML?: string
   projectDir: string
   runIsolate: RunIsolate
   servePath: string
-  importMapFromTOML?: string
 }
 
 /**
@@ -92,6 +94,7 @@ function traverseLocalDependencies(
 export class EdgeFunctionsRegistry {
   public importMapFromDeployConfig?: string
 
+  private aiGatewayContext?: AIGatewayContext | null
   private buildError: Error | null = null
   private bundler: typeof import('@netlify/edge-bundler')
   private configPath: string
@@ -125,6 +128,7 @@ export class EdgeFunctionsRegistry {
   private command: BaseCommand
 
   constructor({
+    aiGatewayContext,
     bundler,
     command,
     config,
@@ -138,6 +142,7 @@ export class EdgeFunctionsRegistry {
     runIsolate,
     servePath,
   }: EdgeFunctionsRegistryOptions) {
+    this.aiGatewayContext = aiGatewayContext
     this.command = command
     this.bundler = bundler
     this.configPath = configPath

src/lib/edge-functions/registry.ts

@@ -1,3 +1,5 @@
+import type { AIGatewayContext } from '@netlify/ai/bootstrap'
+
 import type BaseCommand from '../commands/base-command.js'
 import type { $TSFixMe, NetlifyOptions } from '../commands/types.js'
 import type { BlobsContextWithEdgeAccess } from '../lib/blobs/blobs.js'
@@ -42,6 +44,7 @@ export const generateInspectSettings = (
 export const startProxyServer = async ({
   accountId,
   addonsUrls,
+  aiGatewayContext,
   api,
   blobsContext,
   command,
@@ -65,6 +68,7 @@ export const startProxyServer = async ({
 }: {
   accountId: string | undefined
   addonsUrls: $TSFixMe
+  aiGatewayContext?: AIGatewayContext | null
   api?: NetlifyOptions['api']
   blobsContext?: BlobsContextWithEdgeAccess
   command: BaseCommand
@@ -89,6 +93,7 @@ export const startProxyServer = async ({
 }) => {
   const url = await startProxy({
     addonsUrls,
+    aiGatewayContext,
     blobsContext,
     command,
     config,

src/utils/proxy-server.ts

@@ -14,6 +14,7 @@ import util from 'util'
 import zlib from 'zlib'
 
 import { renderFunctionErrorPage } from '@netlify/dev-utils'
+import type { AIGatewayContext } from '@netlify/ai/bootstrap'
 import contentType from 'content-type'
 import cookie from 'cookie'
 import { getProperty } from 'dot-prop'
@@ -912,6 +913,7 @@ type EdgeFunctionsProxy = Awaited<ReturnType<typeof initializeEdgeFunctionsProxy
 export const startProxy = async function ({
   accountId,
   addonsUrls,
+  aiGatewayContext,
   api,
   blobsContext,
   command,
@@ -937,6 +939,7 @@ export const startProxy = async function ({
   settings: ServerSettings
   disableEdgeFunctions: boolean
   getUpdatedConfig: () => Promise<NormalizedCachedConfigConfig>
+  aiGatewayContext?: AIGatewayContext | null
 } & Record<string, $TSFixMe>) {
   const secondaryServerPort = settings.https ? await getAvailablePort() : null
   const functionsServer = settings.functionsPort ? `http://127.0.0.1:${settings.functionsPort}` : null
@@ -949,8 +952,10 @@ export const startProxy = async function ({
     )
   } else {
     edgeFunctionsProxy = await initializeEdgeFunctionsProxy({
-      command,
+      accountId,
+      aiGatewayContext,
       blobsContext,
+      command,
       config,
       configPath,
       debug,
@@ -966,7 +971,6 @@ export const startProxy = async function ({
       projectDir,
       repositoryRoot,
       siteInfo,
-      accountId,
       state,
     })
   }