feat: support for GCP Context Cache

api/v1alpha1/ai_service_backend.go

@@ -76,10 +76,46 @@ type AIServiceBackendSpec struct {
 	// +optional
 	BodyMutation *HTTPBodyMutation `json:"bodyMutation,omitempty"`
 
+	// ContextCache configures context caching for this backend.
+	// When set, the ai-gateway will automatically resolve or create a context cache
+	// entry for requests that carry Anthropic-style cache_control markers, and inject the
+	// resolved cache resource name into the API request.
+	// Currently only supported for GCP Vertex AI backends.
+	// +optional
+	ContextCache *ContextCacheSpec `json:"contextCache,omitempty"`
+
 	// TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost.
 	// 	That may be useful for the backend that has a different cost calculation logic.
 }
 
+// ContextCacheSpec configures context caching for a backend.
+//
+// Setting this field enables context caching: cache_control markers in requests are
+// resolved against the shared store at URL, and a cache entry is created when absent.
+// Leaving it unset disables caching entirely, and markers are ignored.
+type ContextCacheSpec struct {
+	// DefaultTTL is the default time-to-live for newly created cache entries, expressed
+	// as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used.
+	// Must be at least "60s" per GCP API requirements.
+	//
+	// +optional
+	// +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$`
+	DefaultTTL string `json:"defaultTTL,omitempty"`
+
+	// URL locates the shared Redis store that records which cachedContents entry a given
+	// request prefix resolved to.
+	//
+	// It is given either as a bare "host:port"
+	// (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL.
+	//
+	// If Redis is unreachable at request time, caching is skipped and the request is
+	// served uncached rather than failed.
+	//
+	// +kubebuilder:validation:Required
+	// +kubebuilder:validation:MinLength=1
+	URL string `json:"url"`
+}
+
 // HTTPHeaderMutation defines the mutation of HTTP headers that will be applied to the request
 type HTTPHeaderMutation struct {
 	// Set overwrites/adds the request with the given header (name, value)

api/v1beta1/ai_service_backend.go

@@ -91,6 +91,42 @@ type AIServiceBackendSpec struct {
 	// +kubebuilder:validation:MaxItems=16
 	HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"`
 
+	// ContextCache configures context caching for this backend.
+	// When set, the ai-gateway will automatically resolve or create a context cache
+	// entry for requests that carry Anthropic-style cache_control markers, and inject the
+	// resolved cache resource name into the API request.
+	// Currently only supported for GCP Vertex AI backends.
+	// +optional
+	ContextCache *ContextCacheSpec `json:"contextCache,omitempty"`
+
 	// TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost.
 	// 	That may be useful for the backend that has a different cost calculation logic.
 }
+
+// ContextCacheSpec configures context caching for a backend.
+//
+// Setting this field enables context caching: cache_control markers in requests are
+// resolved against the shared store at URL, and a cache entry is created when absent.
+// Leaving it unset disables caching entirely, and markers are ignored.
+type ContextCacheSpec struct {
+	// DefaultTTL is the default time-to-live for newly created cache entries, expressed
+	// as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used.
+	// Must be at least "60s" per GCP API requirements.
+	//
+	// +optional
+	// +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$`
+	DefaultTTL string `json:"defaultTTL,omitempty"`
+
+	// URL locates the shared Redis store that records which cachedContents entry a given
+	// request prefix resolved to.
+	//
+	// It is given either as a bare "host:port"
+	// (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL.
+	//
+	// If Redis is unreachable at request time, caching is skipped and the request is
+	// served uncached rather than failed.
+	//
+	// +kubebuilder:validation:Required
+	// +kubebuilder:validation:MinLength=1
+	URL string `json:"url"`
+}

cmd/extproc/mainlib/main.go

@@ -425,6 +425,8 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
 	go func() {
 		<-ctx.Done()
 		s.GracefulStop()
+		// Stop context-cache syncers once no request can reach a resolver.
+		server.Close()
 
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.14.0
 	github.com/a8m/envsubst v1.4.3
 	github.com/alecthomas/kong v1.15.0
+	github.com/alicebob/miniredis/v2 v2.33.0
 	github.com/andybalholm/brotli v1.2.2
 	github.com/anthropics/anthropic-sdk-go v1.52.0
 	github.com/aws/aws-sdk-go-v2 v1.42.1
@@ -35,6 +36,7 @@ require (
 	github.com/prometheus/client_golang v1.24.1
 	github.com/prometheus/client_model v0.6.2
 	github.com/prometheus/common v0.70.1
+	github.com/redis/go-redis/v9 v9.7.3
 	github.com/stretchr/testify v1.11.1
 	github.com/testcontainers/testcontainers-go v0.43.0
 	github.com/tetratelabs/func-e v1.6.0
@@ -90,6 +92,7 @@ require (
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.7.2 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/NYTimes/gziphandler v1.1.1 // indirect
+	github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.1 // indirect
 	github.com/avast/retry-go/v5 v5.0.0 // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.28 // indirect
@@ -123,6 +126,7 @@ require (
 	github.com/coreos/go-systemd/v22 v22.7.0 // indirect
 	github.com/cpuguy83/dockercfg v0.3.2 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/cli v29.7.2+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.9.8 // indirect

go.sum

@@ -40,6 +40,10 @@ github.com/alecthomas/kong v1.15.0 h1:BVJstKbpO73zKpmIu+m/aLRrNmWwxXPIGTNin9VmLV
 github.com/alecthomas/kong v1.15.0/go.mod h1:wrlbXem1CWqUV5Vbmss5ISYhsVPkBb1Yo7YKJghju2I=
 github.com/alecthomas/repr v0.5.2 h1:SU73FTI9D1P5UNtvseffFSGmdNci/O6RsqzeXJtP0Qs=
 github.com/alecthomas/repr v0.5.2/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 h1:uvdUDbHQHO85qeSydJtItA4T55Pw6BtAejd0APRJOCE=
+github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302/go.mod h1:SGnFV6hVsYE877CKEZ6tDNTjaSXYUk6QqoIK6PrAtcc=
+github.com/alicebob/miniredis/v2 v2.33.0 h1:uvTF0EDeu9RLnUEG27Db5I68ESoIxTiXbNUiji6lZrA=
+github.com/alicebob/miniredis/v2 v2.33.0/go.mod h1:MhP4a3EU7aENRi9aO+tHfTBZicLqQevyi/DJpoj6mi0=
 github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM=
 github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
 github.com/anthropics/anthropic-sdk-go v1.52.0 h1:1TB9jt4DN87VMwS/hB1VK26tYzK0ipEOtqPaPGFtJQg=
@@ -86,6 +90,10 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=
 github.com/blang/semver/v4 v4.0.0/go.mod h1:IbckMUScFkM3pff0VJDNKRiT6TG/YpiHIM2yvyW5YoQ=
+github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
+github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
 github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk=
 github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
 github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=
@@ -135,6 +143,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
 github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
 github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
@@ -385,6 +395,8 @@ github.com/prometheus/otlptranslator v1.0.0 h1:s0LJW/iN9dkIH+EnhiD3BlkkP5QVIUVEo
 github.com/prometheus/otlptranslator v1.0.0/go.mod h1:vRYWnXvI6aWGpsdY/mOT/cbeVRBlPWtBNDb7kGR3uKM=
 github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI=
 github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY=
+github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
+github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA=
 github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
 github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=

internal/apischema/gcp/gcp.go

@@ -6,6 +6,8 @@
 package gcp
 
 import (
+	"time"
+
 	"google.golang.org/genai"
 
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
@@ -39,6 +41,11 @@ type GenerateContentRequest struct {
 	//
 	// https://github.com/googleapis/go-genai/blob/6a8184fcaf8bf15f0c566616a7b356560309be9b/types.go#L1057
 	SafetySettings []*genai.SafetySetting `json:"safetySettings,omitempty"`
+	// Optional. The name of a pre-existing cached content resource to use as context for generation.
+	// Format: "projects/{project}/locations/{location}/cachedContents/{cache_id}"
+	//
+	// https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/context-cache/context-cache-overview
+	CachedContent string `json:"cachedContent,omitempty"`
 }
 
 // https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#syntax
@@ -171,3 +178,82 @@ type CountTokenRequest struct {
 	// Optional. Configuration that the model uses to generate the response.
 	GenerationConfig *genai.GenerationConfig `json:"generationConfig,omitempty"`
 }
+
+// CreateCachedContent represents the request body for creating a cached content resource in GCP Vertex AI.
+type CreateCachedContent struct {
+	// Required. The model to use for generating the cached content.
+	Model string `json:"model"`
+	// Optional. The user-generated meaningful display name of the cached content.
+	DisplayName string `json:"displayName,omitempty"`
+	// Optional. The TTL for this resource. The expiration time is computed: now + TTL.
+	// Value must be a string in GCP duration format, e.g. "300s".
+	TTL string `json:"ttl,omitempty"`
+	// Optional. Timestamp of when this resource is considered expired.
+	// A pointer so that it is omitted when unset: omitempty has no effect on a
+	// time.Time value, which would otherwise serialize as "0001-01-01T00:00:00Z".
+	ExpireTime *time.Time `json:"expireTime,omitempty"`
+	// The content to cache.
+	Contents []genai.Content `json:"contents"`
+	// Optional. Developer set system instruction.
+	SystemInstruction *genai.Content `json:"systemInstruction,omitempty"`
+	// Optional. A list of `Tools` the model may use to generate the next response.
+	Tools []genai.Tool `json:"tools,omitempty"`
+	// Optional. Configuration for the tools to use. This config is shared for all tools.
+	ToolConfig *genai.ToolConfig `json:"toolConfig,omitempty"`
+	// Optional. The Cloud KMS resource identifier of the customer managed
+	// encryption key used to protect a resource.
+	// The key needs to be in the same region as where the compute resource is
+	// created. See
+	// https://cloud.google.com/vertex-ai/docs/general/cmek for more
+	// details. If this is set, then all created CachedContent objects
+	// will be encrypted with the provided encryption key.
+	// Allowed formats: projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}
+	EncryptionSpec *EncryptionSpec `json:"encryption_spec,omitempty"`
+}
+
+// EncryptionSpec specifies the encryption key that will be used to protect the resource.
+type EncryptionSpec struct {
+	// Required. The Cloud KMS resource identifier of the customer managed
+	// encryption key used to protect a resource.
+	// The key needs to be in the same region as where the compute resource is
+	// created. See
+	// https://cloud.google.com/vertex-ai/docs/general/cmek for more
+	// details. If this is set, then all created CachedContent objects
+	// will be encrypted with the provided encryption key.
+	// Allowed formats: projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}
+	KmsKeyName string `json:"kmsKeyName"`
+}
+
+// A resource used in LLM queries for users to explicitly specify what to cache.
+// https://github.com/googleapis/go-genai/blob/5fa73d012b899ad08135ce9439b88f592acdc5a8/types.go#L6476-L6492
+type CachedContent struct {
+	// Optional. The server-generated resource name of the cached content.
+	Name string `json:"name,omitempty"`
+	// Optional. The user-generated meaningful display name of the cached content.
+	DisplayName string `json:"displayName,omitempty"`
+	// Optional. The name of the publisher model to use for cached content.
+	Model string `json:"model,omitempty"`
+	// Optional. Creation time of the cache entry.
+	CreateTime time.Time `json:"createTime,omitempty"`
+	// Optional. When the cache entry was last updated in UTC time.
+	UpdateTime time.Time `json:"updateTime,omitempty"`
+	// Optional. Expiration time of the cached content.
+	ExpireTime time.Time `json:"expireTime,omitempty"`
+	// Optional. Metadata on the usage of the cached content.
+	UsageMetadata *CachedContentUsageMetadata `json:"usageMetadata,omitempty"`
+}
+
+// Metadata on the usage of the cached content.
+// https://github.com/googleapis/go-genai/blob/5fa73d012b899ad08135ce9439b88f592acdc5a8/types.go#L6462-L6474
+type CachedContentUsageMetadata struct {
+	// Duration of audio in seconds. This field is not supported in Gemini API.
+	AudioDurationSeconds int32 `json:"audioDurationSeconds,omitempty"`
+	// Number of images. This field is not supported in Gemini API.
+	ImageCount int32 `json:"imageCount,omitempty"`
+	// Number of text characters. This field is not supported in Gemini API.
+	TextCount int32 `json:"textCount,omitempty"`
+	// Total number of tokens that the cached content consumes.
+	TotalTokenCount int32 `json:"totalTokenCount,omitempty"`
+	// Duration of video in seconds. This field is not supported in Gemini API.
+	VideoDurationSeconds int32 `json:"videoDurationSeconds,omitempty"`
+}

internal/apischema/openai/openai.go

@@ -2069,6 +2069,15 @@ type GCPVertexAIVendorFields struct {
 	//
 	// https://cloud.google.com/vertex-ai/docs/reference/rest/v1/SafetySetting
 	SafetySettings []*genai.SafetySetting `json:"safetySettings,omitzero"`
+
+	// CachedContent specifies a pre-existing Vertex AI cached content resource to use as context.
+	// Format: "projects/{project}/locations/{location}/cachedContents/{cache_id}"
+	//
+	// When provided, cache_control markers on messages must be absent. The cache reference is
+	// sent directly to Gemini without calling any cache resolution logic.
+	//
+	// https://cloud.google.com/vertex-ai/docs/context-cache/context-cache-overview
+	CachedContent string `json:"cachedContent,omitzero"`
 }
 
 // GCPVertexAIGenerationConfig represents Gemini generation configuration options.

internal/backendauth/gcp.go

@@ -96,3 +96,19 @@ func (g *gcpHandler) Do(_ context.Context, requestHeaders map[string]string, _ [
 	requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", accessToken)
 	return []internalapi.Header{{":path", newPath}, {"Authorization", fmt.Sprintf("Bearer %s", accessToken)}}, nil
 }
+
+// GCPTokenSource implements [filterapi.GCPAuthHandler].
+// It returns the OAuth2 token source used to authenticate to GCP APIs.
+// When a static access token is configured, it wraps it in a StaticTokenSource.
+func (g *gcpHandler) GCPTokenSource() oauth2.TokenSource {
+	if g.tokenSource != nil {
+		return g.tokenSource
+	}
+	return oauth2.StaticTokenSource(&oauth2.Token{AccessToken: g.gcpAccessToken})
+}
+
+// GCPRegion implements [filterapi.GCPAuthHandler].
+func (g *gcpHandler) GCPRegion() string { return g.region }
+
+// GCPProject implements [filterapi.GCPAuthHandler].
+func (g *gcpHandler) GCPProject() string { return g.projectName }