Merge branch 'main' of ssh://github.com/envoyproxy/ai-gateway

Author: siddharth1036

.github/workflows/build_and_test.yaml

@@ -67,32 +67,7 @@ jobs:
             ~/go/pkg/mod
             ~/go/bin
           key: unittest-${{ hashFiles('**/go.mod', '**/go.sum', '**/Makefile') }}-${{ matrix.os }}
-
-      # This runs ollama server to be used in `aigw run` end-to-end tests.
-      # The test case using it will be skipped if ollama is not available.
-      # Since installing it and pulling the model takes a while, we do it only for Linux runners.
-      - name: Start Ollama server
-        if: matrix.os == 'ubuntu-latest'
-        run: |
-          curl -fsSL https://ollama.com/install.sh | sh && sudo systemctl stop ollama
-          nohup ollama serve > ollama.log 2>&1 &
-          timeout 30 sh -c 'until nc -z localhost 11434; do sleep 1; done'
-          grep _MODEL .env.ollama | cut -d= -f2 | xargs -I{} ollama pull {}
-        env:
-          OLLAMA_CONTEXT_LENGTH: 131072 # Larger context for goose
-          OLLAMA_HOST: 0.0.0.0
-      # Download Envoy via func-e using implicit default version `aigw` would
-      # otherwise need to download during test runs.
-      - name: Download Envoy via func-e
-        run: go tool -modfile=tools/go.mod func-e run --version
-        env:
-          FUNC_E_HOME: /tmp/envoy-gateway # hard-coded directory in EG
-      - env:
-          TEST_AWS_ACCESS_KEY_ID: ${{ secrets.AWS_BEDROCK_USER_AWS_ACCESS_KEY_ID }}
-          TEST_AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_BEDROCK_USER_AWS_SECRET_ACCESS_KEY }}
-          TEST_OPENAI_API_KEY: ${{ secrets.ENVOY_AI_GATEWAY_OPENAI_API_KEY }}
-          TEST_GITHUB_ACCESS_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: make test-coverage
+      - run: make test-coverage
       - if: failure()
         run: cat ollama.log || true
       - name: Upload coverage to Codecov
@@ -252,6 +227,7 @@ jobs:
           TEST_AWS_ACCESS_KEY_ID: ${{ secrets.AWS_BEDROCK_USER_AWS_ACCESS_KEY_ID }}
           TEST_AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_BEDROCK_USER_AWS_SECRET_ACCESS_KEY }}
           TEST_OPENAI_API_KEY: ${{ secrets.ENVOY_AI_GATEWAY_OPENAI_API_KEY }}
+          TEST_ANTHROPIC_API_KEY: ${{ secrets.ENVOY_AI_GATEWAY_ANTHROPIC_API_KEY }}
           TEST_GEMINI_API_KEY: ${{ secrets.ENVOY_AI_GATEWAY_GEMINI_API_KEY }}
         run: make test-e2e
 
@@ -388,10 +364,14 @@ jobs:
           FUNC_E_HOME: /tmp/envoy-gateway # hard-coded directory in EG
       - name: Install Goose
         env:
-          GOOSE_VERSION: v1.8.0
+          GOOSE_VERSION: v1.10.0
+          OS: Linux
         run: |
           curl -fsSL https://github.com/block/goose/releases/download/stable/download_cli.sh | CONFIGURE=false bash
-      - run: make test-e2e-aigw
+      - env:
+          # This is used to access the GitHub MCP server.
+          TEST_GITHUB_ACCESS_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: make test-e2e-aigw
       - if: failure()
         run: cat ollama.log || true
 

api/v1alpha1/ai_gateway_route.go

@@ -103,6 +103,8 @@ type AIGatewayRouteSpec struct {
 	//	  type: OutputToken
 	//	- metadataKey: llm_total_token
 	//	  type: TotalToken
+	//	- metadataKey: llm_cached_input_token
+	//	  type: CachedInputToken
 	// ```
 	// Then, with the following BackendTrafficPolicy of Envoy Gateway, you can have three
 	// rate limit buckets for each unique x-user-id header value. One bucket is for the input token,

api/v1alpha1/backendsecurity_policy.go

@@ -19,6 +19,7 @@ const (
 	BackendSecurityPolicyTypeAPIKey           BackendSecurityPolicyType = "APIKey"
 	BackendSecurityPolicyTypeAWSCredentials   BackendSecurityPolicyType = "AWSCredentials"
 	BackendSecurityPolicyTypeAzureAPIKey      BackendSecurityPolicyType = "AzureAPIKey"
+	BackendSecurityPolicyTypeAnthropicAPIKey  BackendSecurityPolicyType = "AnthropicAPIKey" // #nosec G101
 	BackendSecurityPolicyTypeAzureCredentials BackendSecurityPolicyType = "AzureCredentials"
 	BackendSecurityPolicyTypeGCPCredentials   BackendSecurityPolicyType = "GCPCredentials"
 )
@@ -43,11 +44,12 @@ type BackendSecurityPolicy struct {
 //
 // Only one type of BackendSecurityPolicy can be defined.
 // +kubebuilder:validation:MaxProperties=3
-// +kubebuilder:validation:XValidation:rule="self.type == 'APIKey' ? (has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.gcpCredentials)) : true",message="When type is APIKey, only apiKey field should be set"
-// +kubebuilder:validation:XValidation:rule="self.type == 'AWSCredentials' ? (has(self.awsCredentials) && !has(self.apiKey) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.gcpCredentials)) : true",message="When type is AWSCredentials, only awsCredentials field should be set"
-// +kubebuilder:validation:XValidation:rule="self.type == 'AzureAPIKey' ? (has(self.azureAPIKey) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureCredentials) && !has(self.gcpCredentials)) : true",message="When type is AzureAPIKey, only azureAPIKey field should be set"
-// +kubebuilder:validation:XValidation:rule="self.type == 'AzureCredentials' ? (has(self.azureCredentials) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.gcpCredentials)) : true",message="When type is AzureCredentials, only azureCredentials field should be set"
-// +kubebuilder:validation:XValidation:rule="self.type == 'GCPCredentials' ? (has(self.gcpCredentials) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.azureCredentials)) : true",message="When type is GCPCredentials, only gcpCredentials field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'APIKey' ? (has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.gcpCredentials) && !has(self.anthropicAPIKey)) : true",message="When type is APIKey, only apiKey field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'AWSCredentials' ? (has(self.awsCredentials) && !has(self.apiKey) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.gcpCredentials) && !has(self.anthropicAPIKey)) : true",message="When type is AWSCredentials, only awsCredentials field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'AzureAPIKey' ? (has(self.azureAPIKey) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureCredentials) && !has(self.gcpCredentials) && !has(self.anthropicAPIKey)) : true",message="When type is AzureAPIKey, only azureAPIKey field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'AzureCredentials' ? (has(self.azureCredentials) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.gcpCredentials) && !has(self.anthropicAPIKey)) : true",message="When type is AzureCredentials, only azureCredentials field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'GCPCredentials' ? (has(self.gcpCredentials) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.anthropicAPIKey)) : true",message="When type is GCPCredentials, only gcpCredentials field should be set"
+// +kubebuilder:validation:XValidation:rule="self.type == 'AnthropicAPIKey' ? (has(self.anthropicAPIKey) && !has(self.apiKey) && !has(self.awsCredentials) && !has(self.azureAPIKey) && !has(self.azureCredentials) && !has(self.gcpCredentials)) : true",message="When type is AnthropicAPIKey, only anthropicAPIKey field should be set"
 type BackendSecurityPolicySpec struct {
 	// TargetRefs are the names of the AIServiceBackend resources this BackendSecurityPolicy is being attached to.
 	// Attaching multiple BackendSecurityPolicies to the same AIServiceBackend is invalid and will result in an error
@@ -60,7 +62,7 @@ type BackendSecurityPolicySpec struct {
 
 	// Type specifies the type of the backend security policy.
 	//
-	// +kubebuilder:validation:Enum=APIKey;AWSCredentials;AzureAPIKey;AzureCredentials;GCPCredentials
+	// +kubebuilder:validation:Enum=APIKey;AWSCredentials;AzureAPIKey;AzureCredentials;GCPCredentials;AnthropicAPIKey
 	Type BackendSecurityPolicyType `json:"type"`
 
 	// APIKey is a mechanism to access a backend(s). The API key will be injected into the Authorization header.
@@ -82,10 +84,17 @@ type BackendSecurityPolicySpec struct {
 	//
 	// +optional
 	AzureCredentials *BackendSecurityPolicyAzureCredentials `json:"azureCredentials,omitempty"`
+
 	// GCPCredentials is a mechanism to access a backend(s). GCP specific logic will be applied.
 	//
 	// +optional
 	GCPCredentials *BackendSecurityPolicyGCPCredentials `json:"gcpCredentials,omitempty"`
+
+	// AnthropicAPIKey is a mechanism to access Anthropic backend(s). The API key will be injected into the "x-api-key" header.
+	// https://docs.claude.com/en/api/overview#authentication
+	//
+	// +optional
+	AnthropicAPIKey *BackendSecurityPolicyAnthropicAPIKey `json:"anthropicAPIKey,omitempty"`
 }
 
 // BackendSecurityPolicyList contains a list of BackendSecurityPolicy
@@ -314,3 +323,11 @@ type GCPCredentialsFile struct {
 	// The secret should contain the GCP service account credentials file keyed on "service_account.json".
 	SecretRef *gwapiv1.SecretObjectReference `json:"secretRef"`
 }
+
+// BackendSecurityPolicyAnthropicAPIKey specifies the Anthropic API key.
+type BackendSecurityPolicyAnthropicAPIKey struct {
+	// SecretRef is the reference to the secret containing the Anthropic API key.
+	// ai-gateway must be given the permission to read this secret.
+	// The key of the secret should be "apiKey".
+	SecretRef *gwapiv1.SecretObjectReference `json:"secretRef"`
+}

api/v1alpha1/shared_types.go

@@ -15,7 +15,7 @@ package v1alpha1
 type VersionedAPISchema struct {
 	// Name is the name of the API schema of the AIGatewayRoute or AIServiceBackend.
 	//
-	// +kubebuilder:validation:Enum=OpenAI;AWSBedrock;AzureOpenAI;GCPVertexAI;GCPAnthropic
+	// +kubebuilder:validation:Enum=OpenAI;AWSBedrock;AzureOpenAI;GCPVertexAI;GCPAnthropic;Anthropic
 	Name APISchema `json:"name"`
 
 	// Version is the version of the API schema.
@@ -62,6 +62,9 @@ const (
 	//
 	// https://docs.anthropic.com/en/api/claude-on-vertex-ai
 	APISchemaGCPAnthropic APISchema = "GCPAnthropic"
+	// APISchemaAnthropic is the native Anthropic API schema.
+	// https://docs.claude.com/en/home
+	APISchemaAnthropic APISchema = "Anthropic"
 )
 
 const (
@@ -80,7 +83,7 @@ type LLMRequestCost struct {
 	// and it uses "output token" as the cost. The other types are "InputToken", "TotalToken",
 	// and "CEL".
 	//
-	// +kubebuilder:validation:Enum=OutputToken;InputToken;TotalToken;CEL
+	// +kubebuilder:validation:Enum=OutputToken;InputToken;CachedInputToken;TotalToken;CEL
 	Type LLMRequestCostType `json:"type"`
 	// CEL is the CEL expression to calculate the cost of the request.
 	// The CEL expression must return a signed or unsigned integer. If the
@@ -91,13 +94,15 @@ type LLMRequestCost struct {
 	//	* model: the model name extracted from the request content. Type: string.
 	//	* backend: the backend name in the form of "name.namespace". Type: string.
 	//	* input_tokens: the number of input tokens. Type: unsigned integer.
+	//	* cached_input_tokens: the number of cached input tokens. Type: unsigned integer.
 	//	* output_tokens: the number of output tokens. Type: unsigned integer.
 	//	* total_tokens: the total number of tokens. Type: unsigned integer.
 	//
 	// For example, the following expressions are valid:
 	//
 	// 	* "model == 'llama' ?  input_tokens + output_token * 0.5 : total_tokens"
 	//	* "backend == 'foo.default' ?  input_tokens + output_tokens : total_tokens"
+	//	* "backend == 'bar.default' ?  (input_tokens - cached_input_tokens) + cached_input_tokens * 0.1 + output_tokens : total_tokens"
 	//	* "input_tokens + output_tokens + total_tokens"
 	//	* "input_tokens * output_tokens"
 	//
@@ -111,6 +116,8 @@ type LLMRequestCostType string
 const (
 	// LLMRequestCostTypeInputToken is the cost type of the input token.
 	LLMRequestCostTypeInputToken LLMRequestCostType = "InputToken"
+	// LLMRequestCostTypeCachedInputToken is the cost type of the cached input token.
+	LLMRequestCostTypeCachedInputToken LLMRequestCostType = "CachedInputToken"
 	// LLMRequestCostTypeOutputToken is the cost type of the output token.
 	LLMRequestCostTypeOutputToken LLMRequestCostType = "OutputToken"
 	// LLMRequestCostTypeTotalToken is the cost type of the total token.

api/v1alpha1/zz_generated.deepcopy.go

@@ -11,7 +11,7 @@ import (
 	"log/slog"
 	"time"
 
-	"github.com/envoyproxy/ai-gateway/internal/aigw"
+	"github.com/tetratelabs/func-e/experimental/admin"
 )
 
 // healthcheck performs looks up the Envoy subprocess, gets its admin port,
@@ -27,13 +27,12 @@ func healthcheck(ctx context.Context, _, stderr io.Writer) error {
 }
 
 func doHealthcheck(ctx context.Context, aigwPid int, logger *slog.Logger) error {
-	envoyAdmin, err := aigw.NewEnvoyAdminClient(ctx, aigwPid, 0)
-	if err != nil {
+	if adminClient, err := admin.NewAdminClient(ctx, aigwPid); err != nil {
 		logger.Error("Failed to find Envoy admin server", "error", err)
 		return err
-	} else if err = envoyAdmin.IsReady(ctx); err != nil {
-		logger.Error("Envoy admin server is not ready", "adminPort", envoyAdmin.Port(), "error", err)
+	} else if err = adminClient.IsReady(ctx); err != nil {
+		logger.Error("Envoy admin server is not ready", "adminPort", adminClient.Port(), "error", err)
 		return err
 	}
-	return err
+	return nil
 }

cmd/aigw/healthcheck.go

@@ -8,19 +8,16 @@ package main
 import (
 	"bytes"
 	"context"
-	"fmt"
+	"io"
 	"log/slog"
-	"net/http"
-	"net/http/httptest"
-	"net/url"
 	"os"
-	"os/exec"
-	"path/filepath"
-	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
+	func_e "github.com/tetratelabs/func-e"
+	"github.com/tetratelabs/func-e/api"
+	"github.com/tetratelabs/func-e/experimental/admin"
 )
 
 func Test_healthcheck(t *testing.T) {
@@ -38,38 +35,32 @@ func Test_healthcheck(t *testing.T) {
 	})
 
 	t.Run("returns nil when ready", func(t *testing.T) {
-		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			require.Equal(t, "/ready", r.URL.Path)
-			w.WriteHeader(http.StatusOK)
-			_, _ = w.Write([]byte("live"))
-		}))
-		defer server.Close()
-
-		u, err := url.Parse(server.URL)
-		require.NoError(t, err)
-		port, err := strconv.Atoi(u.Port())
-		require.NoError(t, err)
-
-		adminFile := filepath.Join(t.TempDir(), "admin-address.txt")
-		require.NoError(t, os.WriteFile(adminFile, []byte(fmt.Sprintf("127.0.0.1:%d", port)), 0o600))
-
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
 
-		cmdStr := fmt.Sprintf("sleep 30 && echo -- --admin-address-path %s", adminFile)
-		cmd := exec.CommandContext(ctx, "sh", "-c", cmdStr)
-		require.NoError(t, cmd.Start())
-		defer func() {
-			_ = cmd.Process.Kill()
-			_, _ = cmd.Process.Wait()
-		}()
+		var healthCheckErr error
+		var log bytes.Buffer
 
-		time.Sleep(100 * time.Millisecond)
+		// Even though AdminClient.IsReady exists, we don't have it injected in
+		// Docker. This intentionally ignores the parameter.
+		startupHook := func(ctx context.Context, _ admin.AdminClient, _ string) error {
+			logger := slog.New(slog.NewTextHandler(&log, nil))
+			healthCheckErr = doHealthcheck(ctx, pid, logger)
+			// Cancel immediately to stop Envoy and complete test quickly
+			cancel()
+			return nil
+		}
 
-		var buf bytes.Buffer
-		logger := slog.New(slog.NewTextHandler(&buf, nil))
-		err = doHealthcheck(t.Context(), pid, logger)
+		// Run with minimal Envoy config
+		err := func_e.Run(ctx, []string{
+			"--config-yaml",
+			"admin: {address: {socket_address: {address: '127.0.0.1', port_value: 0}}}",
+		}, api.Out(io.Discard), api.EnvoyOut(io.Discard), api.EnvoyErr(io.Discard), admin.WithStartupHook(startupHook))
+
+		// Expect nil error since Run returns nil on context cancellation (documented behavior)
 		require.NoError(t, err)
-		require.Empty(t, buf)
+
+		require.NoError(t, healthCheckErr)
+		require.Empty(t, log)
 	})
 }