diff --git a/cmd/controller/main.go b/cmd/controller/main.go index 67d3f595ba..b455b1de0f 100644 --- a/cmd/controller/main.go +++ b/cmd/controller/main.go @@ -44,6 +44,8 @@ type flags struct { extProcLogFormat string logFormat string extProcEnableRedaction bool + extProcContextCacheStoreURL string + extProcContextCacheDefaultTTL time.Duration extProcImage string extProcImagePullPolicy corev1.PullPolicy enableLeaderElection bool @@ -152,6 +154,16 @@ func parseAndValidateFlags(args []string) (*flags, error) { false, "Enable redaction of sensitive information in debug logs for the external processor.", ) + extProcContextCacheStoreURLPtr := fs.String( + "extProcContextCacheStoreURL", + "", + "Shared store for context caching in the external processor, as host:port or a redis:// URL of a Redis-compatible server. Empty disables context caching.", + ) + extProcContextCacheDefaultTTLPtr := fs.Duration( + "extProcContextCacheDefaultTTL", + 300*time.Second, + "Context cache TTL used when a cache_control marker does not specify one. Whole seconds, at least 60s. Default is 300s.", + ) extProcImagePtr := fs.String( "extProcImage", "docker.io/envoyproxy/ai-gateway-extproc:latest", @@ -379,6 +391,8 @@ func parseAndValidateFlags(args []string) (*flags, error) { extProcLogFormat: *extProcLogFormatPtr, logFormat: *logFormatPtr, extProcEnableRedaction: *extProcEnableRedactionPtr, + extProcContextCacheStoreURL: *extProcContextCacheStoreURLPtr, + extProcContextCacheDefaultTTL: *extProcContextCacheDefaultTTLPtr, extProcImage: *extProcImagePtr, extProcImagePullPolicy: extProcPullPolicy, enableLeaderElection: *enableLeaderElectionPtr, @@ -512,6 +526,8 @@ func main() { ExtProcLogLevel: parsedFlags.extProcLogLevel, ExtProcLogFormat: parsedFlags.extProcLogFormat, ExtProcEnableRedaction: parsedFlags.extProcEnableRedaction, + ExtProcContextCacheStoreURL: parsedFlags.extProcContextCacheStoreURL, + ExtProcContextCacheDefaultTTL: parsedFlags.extProcContextCacheDefaultTTL, EnableLeaderElection: parsedFlags.enableLeaderElection, UDSPath: extProcUDSPath, RequestHeaderAttributes: parsedFlags.requestHeaderAttributes, diff --git a/cmd/extproc/mainlib/main.go b/cmd/extproc/mainlib/main.go index 9830c99b50..00dff78a77 100644 --- a/cmd/extproc/mainlib/main.go +++ b/cmd/extproc/mainlib/main.go @@ -26,6 +26,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/health/grpc_health_v1" + "github.com/envoyproxy/ai-gateway/internal/contextcache/redis" "github.com/envoyproxy/ai-gateway/internal/endpointspec" "github.com/envoyproxy/ai-gateway/internal/extproc" "github.com/envoyproxy/ai-gateway/internal/filterapi" @@ -70,6 +71,11 @@ type extProcFlags struct { maxRecvMsgSize int // endpointPrefixes is the comma-separated key-value pairs for endpoint prefixes. endpointPrefixes string + // contextCacheStoreURL is the shared store for GCP Vertex AI context caching. Empty + // disables context caching. + contextCacheStoreURL string + // contextCacheDefaultTTL is the cache TTL used when a cache_control marker gives none. + contextCacheDefaultTTL time.Duration } // newLogHandler returns the slog handler for the requested output format. Anything other than @@ -162,6 +168,10 @@ func parseAndValidateFlags(args []string) (extProcFlags, error) { "Number of iterations used in the fallback PBKDF2 key derivation for MCP session encryption.") fs.DurationVar(&flags.mcpWriteTimeout, "mcpWriteTimeout", 120*time.Second, "The maximum duration before timing out writes of the MCP response") + fs.StringVar(&flags.contextCacheStoreURL, "contextCacheStoreURL", "", + "Shared store for GCP Vertex AI context caching, as host:port or a redis:// URL of a Redis-compatible server. Empty disables context caching.") + fs.DurationVar(&flags.contextCacheDefaultTTL, "contextCacheDefaultTTL", 300*time.Second, + "Cache TTL used when a cache_control marker does not specify one. Whole seconds, at least 60s.") if err := fs.Parse(args); err != nil { return extProcFlags{}, fmt.Errorf("failed to parse extProcFlags: %w", err) @@ -202,6 +212,11 @@ func parseAndValidateFlags(args []string) (extProcFlags, error) { } } + if flags.contextCacheStoreURL != "" { + if flags.contextCacheDefaultTTL < time.Minute || flags.contextCacheDefaultTTL%time.Second != 0 { + errs = append(errs, fmt.Errorf("contextCacheDefaultTTL must be whole seconds and at least 60s, got %s", flags.contextCacheDefaultTTL)) + } + } return flags, errors.Join(errs...) } @@ -326,6 +341,13 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) { if err != nil { return fmt.Errorf("failed to create external processor server: %w", err) } + if flags.contextCacheStoreURL != "" { + store, storeErr := redis.NewStore(flags.contextCacheStoreURL) + if storeErr != nil { + return fmt.Errorf("invalid contextCacheStoreURL: %w", storeErr) + } + server.EnableContextCache(store, flags.contextCacheDefaultTTL) + } server.Register(path.Join(flags.rootPrefix, endpointPrefixes.OpenAI, "/v1/chat/completions"), extproc.NewFactory( chatCompletionMetricsFactory, tracing.ChatCompletionTracer(), endpointspec.ChatCompletionsEndpointSpec{})) server.Register(path.Join(flags.rootPrefix, endpointPrefixes.OpenAI, "/v1/completions"), extproc.NewFactory( @@ -428,6 +450,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() diff --git a/cmd/extproc/mainlib/main_test.go b/cmd/extproc/mainlib/main_test.go index 95b08a5bfc..dd96031ef6 100644 --- a/cmd/extproc/mainlib/main_test.go +++ b/cmd/extproc/mainlib/main_test.go @@ -385,3 +385,19 @@ backends: require.NoError(t, err, "extproc exited with error before startup message") } } + +func TestParseAndValidateFlags_ContextCache(t *testing.T) { + base := []string{"-configBundlePath", "/dev/null"} + f, err := parseAndValidateFlags(append(base, "-contextCacheStoreURL", "redis:6379")) + require.NoError(t, err) + require.Equal(t, "redis:6379", f.contextCacheStoreURL) + require.Equal(t, 300*time.Second, f.contextCacheDefaultTTL) + + for _, bad := range []string{"30s", "90500ms"} { + _, err = parseAndValidateFlags(append(base, "-contextCacheStoreURL", "redis:6379", "-contextCacheDefaultTTL", bad)) + require.ErrorContains(t, err, "contextCacheDefaultTTL", bad) + } + // The TTL is not validated when caching is off. + _, err = parseAndValidateFlags(append(base, "-contextCacheDefaultTTL", "30s")) + require.NoError(t, err) +} diff --git a/go.mod b/go.mod index 63d191a748..1a07bb9648 100644 --- a/go.mod +++ b/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.45.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.12.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.20.1 // 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.8.0+incompatible // indirect github.com/docker/docker-credential-helpers v0.9.8 // indirect diff --git a/go.sum b/go.sum index 8e7ed239ea..a61ecd4932 100644 --- a/go.sum +++ b/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= @@ -379,6 +389,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.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= diff --git a/internal/apischema/gcp/gcp.go b/internal/apischema/gcp/gcp.go index d2bff89f62..b8b2a62743 100644 --- a/internal/apischema/gcp/gcp.go +++ b/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"` +} diff --git a/internal/backendauth/gcp.go b/internal/backendauth/gcp.go index 4dd77c78f1..0b429494c3 100644 --- a/internal/backendauth/gcp.go +++ b/internal/backendauth/gcp.go @@ -32,7 +32,7 @@ type gcpHandler struct { projectName string // The GCP project to use for requests. } -func newGCPHandler(ctx context.Context, gcpAuth *filterapi.GCPAuth) (filterapi.BackendAuthHandler, error) { +func newGCPHandler(_ context.Context, gcpAuth *filterapi.GCPAuth) (filterapi.BackendAuthHandler, error) { if gcpAuth == nil { return nil, fmt.Errorf("GCP auth configuration cannot be nil") } @@ -47,8 +47,12 @@ func newGCPHandler(ctx context.Context, gcpAuth *filterapi.GCPAuth) (filterapi.B handler.gcpAccessToken = gcpAuth.AccessToken } else { // Use ADC for GKE Workload Identity. TokenSource auto-refreshes in Do(). - // Inject HTTP client with proxy support into context for token operations. - ctx = context.WithValue(ctx, oauth2.HTTPClient, gcpHTTPClient) + // + // The token source keeps this context and uses it for every later token fetch, so + // it must live as long as the handler. The caller's ctx does not: config reloads + // cancel theirs as soon as the reload returns, which would make every fetch fail + // with "context canceled". The context only carries the proxy-aware HTTP client. + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, gcpHTTPClient) creds, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform") if err != nil { return nil, fmt.Errorf("failed to find GCP default credentials: %w", err) @@ -96,3 +100,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 } diff --git a/internal/contextcache/gcpcache/gcpcache.go b/internal/contextcache/gcpcache/gcpcache.go new file mode 100644 index 0000000000..8cd8c0cb83 --- /dev/null +++ b/internal/contextcache/gcpcache/gcpcache.go @@ -0,0 +1,556 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +// Package gcpcache implements context-cache resolution for GCP Vertex AI (Gemini). +// +// When a client sends an OpenAI-format request with Anthropic-style cache_control markers, +// the resolver: +// 1. Finds the last cache_control breakpoint in the message list. +// 2. Splits the request into a cached prefix (tools + system + messages up to the breakpoint) +// and a non-cached remainder. +// 3. Generates a deterministic SHA-256 cache key and looks it up in the shared store. +// 4. On a store miss, creates the cache entry, publishes it to the store, and returns +// the cache resource name and the non-cached messages. +package gcpcache + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "log/slog" + "net/http" + "strconv" + "sync" + "time" + + "google.golang.org/genai" + + "github.com/envoyproxy/ai-gateway/internal/apischema/gcp" + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/json" + "github.com/envoyproxy/ai-gateway/internal/translator" +) + +const ( + // DefaultTTL is the cache TTL used when New is given none and the cache_control marker + // does not specify one. + DefaultTTL = 300 * time.Second + + // gcpCachedContentsBasePath is the base URL for the Vertex AI cachedContents REST API. + gcpCachedContentsBasePath = "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/cachedContents" +) + +// Resolver is the GCP Vertex AI contextcache.CacheResolver and contextcache.CacheSyncer. +type Resolver struct { + httpClient *http.Client + + // store holds resolved cache names. Backed by Redis it is shared across replicas, so + // a name published by one is reused by all. It does not prevent simultaneous creates + // on a cold prefix. Its failures are never fatal: they are logged and treated as a + // miss. + store contextcache.Store + + // logger records store failures, which are otherwise invisible because the request + // proceeds normally. + logger *slog.Logger + + // defaultTTL is the GCP duration string (e.g. "300s") used when a cache_control + // marker does not specify a TTL. + defaultTTL string + + // mu guards auth and the background syncer's lifecycle. + mu sync.Mutex + // auth is the backend's current credentials. Resolve uses them to create caches and + // the background syncer uses them to list them. Nil makes the resolver inert. It is + // replaced by SetAuth on every config reload, so rotated tokens and changed + // project/region take effect without rebuilding the resolver. + auth filterapi.GCPAuthHandler + // cancel stops the syncer; nil until Start runs one. + cancel context.CancelFunc + // done is closed when the syncer goroutine exits. + done chan struct{} + // closed is set by Close, after which Start does nothing. + closed bool +} + +// syncInterval is how often the background syncer runs a round. Between rounds, +// a cache the store has forgotten is not rediscovered, so this bounds how long drift +// persists. It is sized against the 300s default cache TTL. +const syncInterval = 60 * time.Second + +var _ contextcache.CacheResolver = (*Resolver)(nil) + +var _ contextcache.CacheSyncer = (*Resolver)(nil) + +// Start implements contextcache.CacheSyncer. It does nothing if the resolver was built +// without credentials. +func (r *Resolver) Start() { + r.mu.Lock() + defer r.mu.Unlock() + if r.closed || r.cancel != nil || r.auth == nil { + return + } + // A nil auth makes the syncer read the resolver's current credentials each round. + rc := r.newSyncer(nil, syncInterval) + ctx, cancel := context.WithCancel(context.Background()) + r.cancel = cancel + r.done = make(chan struct{}) + go func() { + defer close(r.done) + rc.Run(ctx) + }() +} + +// SetAuth replaces the credentials used by Resolve and the background syncer. LoadConfig +// calls it on every reload, because the auth handler is rebuilt each time (the access +// token is rotated by the controller) while the resolver is reused. +func (r *Resolver) SetAuth(auth filterapi.GCPAuthHandler) { + r.mu.Lock() + defer r.mu.Unlock() + r.auth = auth +} + +// GetAuth returns the credentials in effect. It takes the lock because SetAuth may run +// concurrently from a config reload. +func (r *Resolver) GetAuth() filterapi.GCPAuthHandler { + r.mu.Lock() + defer r.mu.Unlock() + return r.auth +} + +// Close implements contextcache.CacheSyncer. +func (r *Resolver) Close() error { + r.mu.Lock() + r.closed = true + cancel, done := r.cancel, r.done + r.mu.Unlock() + if cancel != nil { + cancel() + <-done + } + return nil +} + +// New creates a new Resolver. +// +// httpClient is used for calls to the GCP cachedContents API; pass nil for a default. +// store holds resolved cache names; pass nil to disable caching entirely, which makes +// resolution inert rather than failing requests. auth is the backend's credentials, used +// for both resolving and syncing; nil makes the resolver inert. logger may be nil. +func New(httpClient *http.Client, store contextcache.Store, auth filterapi.GCPAuthHandler, defaultTTL time.Duration, logger *slog.Logger) *Resolver { + if httpClient == nil { + httpClient = &http.Client{Timeout: 30 * time.Second} + } + if store == nil { + store = contextcache.NoopStore{} + } + if logger == nil { + logger = slog.New(slog.DiscardHandler) + } + if defaultTTL <= 0 { + defaultTTL = DefaultTTL + } + return &Resolver{ + httpClient: httpClient, store: store, logger: logger, auth: auth, + defaultTTL: gcpDuration(defaultTTL), + } +} + +// Resolve implements contextcache.CacheResolver. +func (r *Resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest) (*contextcache.ResolveResult, error) { + // Without credentials nothing can be created, so caching is inert. + auth := r.GetAuth() + if auth == nil { + return nil, nil + } + // Find the last cache_control breakpoint in the message list. + breakpoint := findBreakpoint(openAIReq.Messages) + if breakpoint < 0 { + // No markers — nothing to cache. + return nil, nil + } + + // Split cached prefix. + cachedMessages := openAIReq.Messages[:breakpoint+1] + remainderMessages := openAIReq.Messages[breakpoint+1:] + + // Extract the TTL from the breakpoint message. If not found, fall back to default. + ttl := extractTTL(openAIReq.Messages[breakpoint], r.defaultTTL) + + // Build the Gemini cached prefix (tools + system + contents). + contents, systemInstruction, err := translator.OpenAIMessagesToGeminiContents(cachedMessages, openAIReq.Model) + if err != nil { + return nil, fmt.Errorf("gcpcache: failed to convert cached messages to Gemini format: %w", err) + } + geminiTools, err := translator.OpenAIToolsToGeminiTools(openAIReq.Tools, false) + if err != nil { + return nil, fmt.Errorf("gcpcache: failed to convert tools to Gemini format: %w", err) + } + + // Compute a deterministic cache key. + cacheKey, err := computeCacheKey(auth.GCPProject(), auth.GCPRegion(), openAIReq.Model, contents, systemInstruction, geminiTools) + if err != nil { + return nil, fmt.Errorf("gcpcache: failed to compute cache key: %w", err) + } + + // Check the shared store first. A store failure is not fatal: it is logged and + // treated as a miss, so an unreachable store degrades caching rather than the + // request. GCP-side failures below are a different matter and do fail fast. + if e, ok := r.storeGet(ctx, cacheKey); ok { + r.logger.Debug("gcpcache: cache hit", + slog.String("key", cacheKey), slog.String("cache", e.Name), + slog.Time("expireTime", e.ExpireTime), slog.Int("cachedMessages", breakpoint+1), + slog.Int("remainingMessages", len(remainderMessages))) + return &contextcache.ResolveResult{ + CacheName: e.Name, + Messages: remainderMessages, + ExpireTime: e.ExpireTime, + }, nil + } + + // Store miss — create against GCP. See resolveUncached for how concurrent + // resolutions of the same key behave. + r.logger.Debug("gcpcache: cache miss, creating", + slog.String("key", cacheKey), slog.String("model", openAIReq.Model), + slog.String("ttl", ttl), slog.Int("cachedMessages", breakpoint+1)) + result, err := r.resolveUncached(ctx, openAIReq, auth, cacheKey, ttl, contents, systemInstruction, geminiTools) + if err != nil { + return nil, err + } + r.logger.Debug("gcpcache: cache created", + slog.String("key", cacheKey), slog.String("cache", result.CacheName), + slog.Time("expireTime", result.ExpireTime), slog.Int("tokenCount", result.TokenCount), + slog.Int("remainingMessages", len(remainderMessages))) + + // resolveUncached leaves Messages unset because it deals only in the cached prefix; + // the remainder is per-request. + result.Messages = remainderMessages + return result, nil +} + +// resolveUncached creates a GCP cachedContents entry for a cache key that missed the +// store, and publishes it. It deals only in the cached prefix, so it leaves +// ResolveResult.Messages unset for the caller to fill in per-request. +// +// It neither waits on other replicas nor lists existing caches. Replicas that miss the +// same cold prefix at the same moment each create; each keeps and bills its own cache +// (Created is true, so the write shows up in cache-write token metrics), and the surplus +// expires with its TTL. +func (r *Resolver) resolveUncached( + ctx context.Context, + openAIReq *openai.ChatCompletionRequest, + gcpAuth filterapi.GCPAuthHandler, + cacheKey, ttl string, + contents []genai.Content, + systemInstruction *genai.Content, + geminiTools []genai.Tool, +) (*contextcache.ResolveResult, error) { + region := gcpAuth.GCPRegion() + project := gcpAuth.GCPProject() + + tokenSrc := gcpAuth.GCPTokenSource() + token, err := tokenSrc.Token() + if err != nil { + return nil, fmt.Errorf("gcpcache: failed to get GCP access token: %w", err) + } + accessToken := token.AccessToken + + // Create without waiting on other replicas. A replica racing on the same cold prefix + // may create too; each keeps the cache it created, and the store holds whichever name + // was published last. Both are valid GCP caches. + baseURL := fmt.Sprintf(gcpCachedContentsBasePath, region, project, region) + created, tokenCount, expireTime, err := r.createCache(ctx, baseURL, accessToken, openAIReq.Model, region, project, cacheKey, contents, systemInstruction, geminiTools, ttl) + if err != nil { + return nil, fmt.Errorf("gcpcache: failed to create cached content: %w", err) + } + + r.storeSet(ctx, cacheKey, contextcache.Entry{Name: created, ExpireTime: expireTime}) + return &contextcache.ResolveResult{ + CacheName: created, + Created: true, + TokenCount: tokenCount, + ExpireTime: expireTime, + }, nil +} + +// ----------------------------------------------------------------------- +// Store access +// +// The resolver never propagates a store failure to its caller: caching is an +// optimization, and an unreachable store must not turn a servable request into an +// error. Failures are logged and treated as a miss, leaving GCP as the source of +// truth. GCP-side failures are propagated, because those are actionable by the +// user — a create rejected for being below the model's minimum token count means the +// cache_control markers are misplaced, and silently serving the request uncached +// would hide a cost increase. +// ----------------------------------------------------------------------- + +// storeGet reads a resolved cache name, reporting a miss on any failure. +func (r *Resolver) storeGet(ctx context.Context, key string) (contextcache.Entry, bool) { + e, ok, err := r.store.Get(ctx, key) + if err != nil { + r.logger.Warn("gcpcache: cache store read failed, proceeding uncached", + slog.String("error", err.Error())) + return contextcache.Entry{}, false + } + if !ok { + return contextcache.Entry{}, false + } + // Treat entries expiring within 10s as stale so there is time to re-resolve before + // the cache disappears underneath an in-flight request. + if time.Until(e.ExpireTime) < contextcache.StaleThreshold { + return contextcache.Entry{}, false + } + return e, true +} + +// storeSet records a resolved cache name, expiring it with the GCP entry itself so +// the store cannot outlive what it points at. +func (r *Resolver) storeSet(ctx context.Context, key string, e contextcache.Entry) { + ttl := time.Until(e.ExpireTime) + if ttl <= 0 { + return + } + if err := r.store.Set(ctx, key, e, ttl); err != nil { + r.logger.Warn("gcpcache: cache store write failed", + slog.String("error", err.Error())) + } +} + +// ----------------------------------------------------------------------- +// Breakpoint and TTL helpers +// ----------------------------------------------------------------------- + +// findBreakpoint returns the index of the last message that contains a +// cache_control marker, or -1 if none is found. +func findBreakpoint(messages []openai.ChatCompletionMessageParamUnion) int { + for i := len(messages) - 1; i >= 0; i-- { + if messageHasCacheControl(&messages[i]) { + return i + } + } + return -1 +} + +// messageHasCacheControl returns true if the message contains at least one +// content part with an Anthropic-style cache_control marker. +func messageHasCacheControl(msg *openai.ChatCompletionMessageParamUnion) bool { + if msg.OfTool != nil && isCacheControlSet(msg.OfTool.AnthropicContentFields) { + return true + } + if msg.OfSystem != nil { + if parts, ok := msg.OfSystem.Content.Value.([]openai.ChatCompletionContentPartTextParam); ok { + for i := range parts { + if isCacheControlSet(parts[i].AnthropicContentFields) { + return true + } + } + } + } + if msg.OfUser != nil { + if parts, ok := msg.OfUser.Content.Value.([]openai.ChatCompletionContentPartUserUnionParam); ok { + for i := range parts { + p := &parts[i] + if p.OfText != nil && isCacheControlSet(p.OfText.AnthropicContentFields) { + return true + } + if p.OfImageURL != nil && isCacheControlSet(p.OfImageURL.AnthropicContentFields) { + return true + } + if p.OfInputAudio != nil && isCacheControlSet(p.OfInputAudio.AnthropicContentFields) { + return true + } + } + } + } + return false +} + +// isCacheControlSet returns true if the AnthropicContentFields carries an ephemeral cache marker. +func isCacheControlSet(fields *openai.AnthropicContentFields) bool { + return fields != nil && string(fields.CacheControl.Type) == "ephemeral" +} + +// extractTTL returns the GCP TTL string (e.g. "3600s") from the last cache_control +// marker found in the message. Falls back to defaultTTL when none is specified. +// Anthropic TTL values ("5m", "1h") are converted to GCP seconds format. +func extractTTL(msg openai.ChatCompletionMessageParamUnion, defaultTTL string) string { + var ttl string + checkFields := func(f *openai.AnthropicContentFields) { + if f != nil && string(f.CacheControl.Type) == "ephemeral" && string(f.CacheControl.TTL) != "" { + ttl = anthropicTTLToGCP(string(f.CacheControl.TTL), defaultTTL) + } + } + if msg.OfTool != nil { + checkFields(msg.OfTool.AnthropicContentFields) + } + if msg.OfSystem != nil { + if parts, ok := msg.OfSystem.Content.Value.([]openai.ChatCompletionContentPartTextParam); ok { + for i := range parts { + checkFields(parts[i].AnthropicContentFields) + } + } + } + if msg.OfUser != nil { + if parts, ok := msg.OfUser.Content.Value.([]openai.ChatCompletionContentPartUserUnionParam); ok { + for i := range parts { + p := &parts[i] + if p.OfText != nil { + checkFields(p.OfText.AnthropicContentFields) + } + if p.OfImageURL != nil { + checkFields(p.OfImageURL.AnthropicContentFields) + } + if p.OfInputAudio != nil { + checkFields(p.OfInputAudio.AnthropicContentFields) + } + } + } + } + if ttl == "" { + return defaultTTL + } + return ttl +} + +// anthropicTTLToGCP converts Anthropic cache TTL values to GCP seconds format. +// Anthropic supports "5m" and "1h"; unknown values fall back to defaultTTL. +func anthropicTTLToGCP(ttl, defaultTTL string) string { + switch ttl { + case "5m": + return "300s" + case "1h": + return "3600s" + default: + // If the value is already in GCP seconds format (e.g. "600s"), pass it through. + if len(ttl) > 1 && ttl[len(ttl)-1] == 's' { + return ttl + } + return defaultTTL + } +} + +// gcpDuration formats d as a GCP duration string in whole seconds, e.g. "300s". +func gcpDuration(d time.Duration) string { + return strconv.FormatInt(int64(d/time.Second), 10) + "s" +} + +// computeKeyInputs converts a cached message prefix into Gemini format for key generation. +func computeKeyInputs(model string, cachedMessages []openai.ChatCompletionMessageParamUnion) ([]genai.Content, *genai.Content, error) { + contents, sys, err := translator.OpenAIMessagesToGeminiContents(cachedMessages, model) + if err != nil { + return nil, nil, err + } + return contents, sys, nil +} + +// ----------------------------------------------------------------------- +// Cache key +// ----------------------------------------------------------------------- + +type cacheKeyInput struct { + // Project and Region scope the key to where the cache lives: a cachedContents name is + // only usable in its own project and region, and backends in different locations may + // share one Redis. + Project string `json:"project"` + Region string `json:"region"` + Model string `json:"model"` + Contents []genai.Content `json:"contents"` + SystemInstruction *genai.Content `json:"systemInstruction,omitempty"` + Tools []genai.Tool `json:"tools,omitempty"` +} + +// computeCacheKey generates a deterministic SHA-256 hex digest over the +// (project, region, model, contents, systemInstruction, tools) tuple. The digest is used +// as the GCP cachedContents displayName. +func computeCacheKey(project, region, model string, contents []genai.Content, systemInstruction *genai.Content, tools []genai.Tool) (string, error) { + input := cacheKeyInput{ + Project: project, + Region: region, + Model: model, + Contents: contents, + SystemInstruction: systemInstruction, + Tools: tools, + } + // Sorted map keys: tool schemas and function call args/results are maps, and an + // unsorted encoding would give a different key for the same prefix on each request. + b, err := json.MarshalSorted(input) + if err != nil { + return "", fmt.Errorf("failed to marshal cache key input: %w", err) + } + h := sha256.Sum256(b) + return hex.EncodeToString(h[:]), nil +} + +// ----------------------------------------------------------------------- +// GCP cachedContents REST API calls +// ----------------------------------------------------------------------- + +// createCache creates a cached content in GCP and returns the new cache's resource name, +// token count, and expiry time. +func (r *Resolver) createCache( + ctx context.Context, + baseURL, accessToken, model, region, project, cacheKey string, + contents []genai.Content, + systemInstruction *genai.Content, + tools []genai.Tool, + ttl string, +) (name string, tokenCount int, expireTime time.Time, err error) { + // Vertex AI expects the full model resource name. + fullModel := fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s", project, region, model) + + body := gcp.CreateCachedContent{ + Model: fullModel, + Contents: contents, + SystemInstruction: systemInstruction, + Tools: tools, + DisplayName: cacheKey, + TTL: ttl, + } + bodyBytes, err := json.Marshal(body) + if err != nil { + return "", 0, time.Time{}, fmt.Errorf("failed to marshal create cache request body: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes)) + if err != nil { + return "", 0, time.Time{}, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Content-Type", "application/json") + + resp, err := r.httpClient.Do(req) + if err != nil { + return "", 0, time.Time{}, err + } + defer resp.Body.Close() + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return "", 0, time.Time{}, err + } + if resp.StatusCode != http.StatusOK { + return "", 0, time.Time{}, fmt.Errorf("create cachedContent returned HTTP %d: %s", resp.StatusCode, respBody) + } + + var cr gcp.CachedContent + if err = json.Unmarshal(respBody, &cr); err != nil { + return "", 0, time.Time{}, fmt.Errorf("failed to decode create cache response: %w", err) + } + // A zero expiry would make storeSet skip the write, so every later request for this + // prefix would create again with nothing logged. Fail visibly instead. + if cr.ExpireTime.IsZero() { + return "", 0, time.Time{}, fmt.Errorf("create cachedContent response has no expireTime: %s", respBody) + } + + if cr.UsageMetadata != nil { + tokenCount = int(cr.UsageMetadata.TotalTokenCount) + } + return cr.Name, tokenCount, cr.ExpireTime, nil +} diff --git a/internal/contextcache/gcpcache/gcpcache_test.go b/internal/contextcache/gcpcache/gcpcache_test.go new file mode 100644 index 0000000000..2b6ee8abd0 --- /dev/null +++ b/internal/contextcache/gcpcache/gcpcache_test.go @@ -0,0 +1,787 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package gcpcache + +import ( + "bytes" + "context" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + anthropic "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/shared/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" + "google.golang.org/genai" + + "github.com/envoyproxy/ai-gateway/internal/apischema/gcp" + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/internalapi" + "github.com/envoyproxy/ai-gateway/internal/json" +) + +// ----------------------------------------------------------------------- +// Fake GCPAuthHandler +// ----------------------------------------------------------------------- + +type fakeGCPAuth struct { + token string + region string + project string +} + +func (f *fakeGCPAuth) Do(_ context.Context, _ map[string]string, _ []byte) ([]internalapi.Header, error) { + return nil, nil +} + +func (f *fakeGCPAuth) GCPTokenSource() oauth2.TokenSource { + return oauth2.StaticTokenSource(&oauth2.Token{AccessToken: f.token}) +} +func (f *fakeGCPAuth) GCPRegion() string { return f.region } +func (f *fakeGCPAuth) GCPProject() string { return f.project } + +var _ filterapi.GCPAuthHandler = (*fakeGCPAuth)(nil) + +// ----------------------------------------------------------------------- +// Message builders +// ----------------------------------------------------------------------- + +func ephemeralFields() *openai.AnthropicContentFields { + return &openai.AnthropicContentFields{ + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + } +} + +func ephemeralFieldsWithTTL(ttl anthropic.CacheControlEphemeralTTL) *openai.AnthropicContentFields { + return &openai.AnthropicContentFields{ + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + TTL: ttl, + }, + } +} + +func systemMsg(text string, fields *openai.AnthropicContentFields) openai.ChatCompletionMessageParamUnion { + return openai.ChatCompletionMessageParamUnion{ + OfSystem: &openai.ChatCompletionSystemMessageParam{ + Role: openai.ChatMessageRoleSystem, + Content: openai.ContentUnion{Value: []openai.ChatCompletionContentPartTextParam{ + {Type: "text", Text: text, AnthropicContentFields: fields}, + }}, + }, + } +} + +func userMsg(text string) openai.ChatCompletionMessageParamUnion { + return openai.ChatCompletionMessageParamUnion{ + OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{Value: text}, + }, + } +} + +// ----------------------------------------------------------------------- +// Fake cachedContents server +// ----------------------------------------------------------------------- + +type fakeCacheServer struct { + srv *httptest.Server + listResponse string + createStatus int + createBody string + // Counters are atomic because the httptest handler runs on a goroutine per + // request, and the concurrency tests below drive many at once. + listCalls atomic.Int64 + createCalls atomic.Int64 +} + +func (f *fakeCacheServer) lists() int { return int(f.listCalls.Load()) } +func (f *fakeCacheServer) creates() int { return int(f.createCalls.Load()) } + +func newFakeCacheServer(t *testing.T, listBody, createBody string, createStatus int) *fakeCacheServer { + t.Helper() + f := &fakeCacheServer{ + listResponse: listBody, + createStatus: createStatus, + createBody: createBody, + } + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.Method { + case http.MethodGet: + f.listCalls.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(f.listResponse)) + case http.MethodPost: + f.createCalls.Add(1) + w.WriteHeader(f.createStatus) + _, _ = w.Write([]byte(f.createBody)) + } + })) + t.Cleanup(f.srv.Close) + return f +} + +// redirectTransport rewrites all outbound requests to hit the fake server host. +type redirectTransport struct { + fakeHost string +} + +func (t *redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) { + req2 := req.Clone(req.Context()) + req2.URL.Scheme = "http" + req2.URL.Host = t.fakeHost + return http.DefaultTransport.RoundTrip(req2) +} + +// resolverWithServer builds a resolver pointed at a fake Google server. An optional store +// may be supplied; with none, the no-op store applies and every lookup misses. +func resolverWithServer(srvURL string, store ...contextcache.Store) *Resolver { + host := srvURL[len("http://"):] + var s contextcache.Store + if len(store) > 0 { + s = store[0] + } + return New(&http.Client{ + Transport: &redirectTransport{fakeHost: host}, + Timeout: 5 * time.Second, + }, s, syncAuth, 0, nil) +} + +// memStore is an in-process contextcache.Store for tests. +type memStore struct { + mu sync.Mutex + entries map[string]contextcache.Entry + // getErr, when set, is returned from every Get, standing in for an unreachable store. + getErr error + // setErr, when set, is returned from every Set. + setErr error + gets atomic.Int64 + sets atomic.Int64 +} + +func newMemStore() *memStore { return &memStore{entries: map[string]contextcache.Entry{}} } + +func (m *memStore) Get(_ context.Context, key string) (contextcache.Entry, bool, error) { + m.gets.Add(1) + if m.getErr != nil { + return contextcache.Entry{}, false, m.getErr + } + m.mu.Lock() + defer m.mu.Unlock() + e, ok := m.entries[key] + return e, ok, nil +} + +func (m *memStore) Set(_ context.Context, key string, e contextcache.Entry, _ time.Duration) error { + m.sets.Add(1) + if m.setErr != nil { + return m.setErr + } + m.mu.Lock() + defer m.mu.Unlock() + m.entries[key] = e + return nil +} + +func (m *memStore) SetNX(_ context.Context, key string, e contextcache.Entry, _ time.Duration) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.entries[key]; ok { + return false, nil + } + m.entries[key] = e + return true, nil +} + +// AcquireGate never wins: resolver tests using memStore do not exercise syncing. +func (m *memStore) AcquireGate(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (m *memStore) seed(key, cacheName string, expireTime time.Time) { + m.mu.Lock() + defer m.mu.Unlock() + m.entries[key] = contextcache.Entry{Name: cacheName, ExpireTime: expireTime} +} + +// computeKeyForRequest replicates the key generation for use in test assertions. +func computeKeyForRequest(t *testing.T, req *openai.ChatCompletionRequest) string { + t.Helper() + bp := findBreakpoint(req.Messages) + require.GreaterOrEqual(t, bp, 0) + contents, sys, err := computeKeyInputs(req.Model, req.Messages[:bp+1]) + require.NoError(t, err) + key, err := computeCacheKey(syncAuth.project, syncAuth.region, req.Model, contents, sys, nil) + require.NoError(t, err) + return key +} + +// ----------------------------------------------------------------------- +// Unit tests: findBreakpoint +// ----------------------------------------------------------------------- + +func TestFindBreakpoint(t *testing.T) { + tests := []struct { + name string + messages []openai.ChatCompletionMessageParamUnion + want int + }{ + { + name: "no markers returns -1", + messages: []openai.ChatCompletionMessageParamUnion{userMsg("hello")}, + want: -1, + }, + { + name: "system marker at index 0", + messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("hello"), + }, + want: 0, + }, + { + name: "last marker wins when multiple exist", + messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q1"), + {OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{ + Value: []openai.ChatCompletionContentPartUserUnionParam{ + {OfText: &openai.ChatCompletionContentPartTextParam{ + Type: "text", + Text: "q2", + AnthropicContentFields: ephemeralFields(), + }}, + }, + }, + }}, + userMsg("q3"), + }, + want: 2, + }, + { + name: "tool message with cache_control", + messages: []openai.ChatCompletionMessageParamUnion{ + {OfTool: &openai.ChatCompletionToolMessageParam{ + Role: openai.ChatMessageRoleTool, + ToolCallID: "call_1", + Content: openai.ContentUnion{Value: "result"}, + AnthropicContentFields: ephemeralFields(), + }}, + userMsg("follow-up"), + }, + want: 0, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, findBreakpoint(tc.messages)) + }) + } +} + +// ----------------------------------------------------------------------- +// Unit tests: extractTTL / anthropicTTLToGCP +// ----------------------------------------------------------------------- + +func TestExtractTTL(t *testing.T) { + tests := []struct { + name string + msg openai.ChatCompletionMessageParamUnion + wantTTL string + }{ + { + name: "no ttl → default 300s", + msg: systemMsg("sys", ephemeralFields()), + wantTTL: "300s", + }, + { + name: "5m → 300s", + msg: systemMsg("sys", ephemeralFieldsWithTTL("5m")), + wantTTL: "300s", + }, + { + name: "1h → 3600s", + msg: systemMsg("sys", ephemeralFieldsWithTTL("1h")), + wantTTL: "3600s", + }, + { + name: "raw GCP seconds passthrough", + msg: systemMsg("sys", ephemeralFieldsWithTTL("600s")), + wantTTL: "600s", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.wantTTL, extractTTL(tc.msg, "300s")) + }) + } +} + +// ----------------------------------------------------------------------- +// Unit tests: computeCacheKey +// ----------------------------------------------------------------------- + +func TestComputeCacheKey_Deterministic(t *testing.T) { + msgs := []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + } + contents, sys, err := computeKeyInputs("gemini-1.5-pro", msgs) + require.NoError(t, err) + + k1, err := computeCacheKey("p", "r", "gemini-1.5-pro", contents, sys, nil) + require.NoError(t, err) + k2, err := computeCacheKey("p", "r", "gemini-1.5-pro", contents, sys, nil) + require.NoError(t, err) + assert.Equal(t, k1, k2) + assert.Len(t, k1, 64, "expected SHA-256 hex digest") +} + +func TestComputeCacheKey_DifferentModels_DifferentKeys(t *testing.T) { + msgs := []openai.ChatCompletionMessageParamUnion{systemMsg("sys", ephemeralFields())} + c, s, err := computeKeyInputs("gemini-1.5-pro", msgs) + require.NoError(t, err) + k1, _ := computeCacheKey("p", "r", "gemini-1.5-pro", c, s, nil) + k2, _ := computeCacheKey("p", "r", "gemini-2.0-flash", c, s, nil) + assert.NotEqual(t, k1, k2) +} + +// A cache name is only usable in its own project and region, so the same prompt must +// key differently per location when backends share a Redis. +func TestComputeCacheKey_DifferentLocations_DifferentKeys(t *testing.T) { + msgs := []openai.ChatCompletionMessageParamUnion{systemMsg("sys", ephemeralFields())} + c, s, err := computeKeyInputs("gemini-1.5-pro", msgs) + require.NoError(t, err) + base, _ := computeCacheKey("p", "us-central1", "gemini-1.5-pro", c, s, nil) + otherRegion, _ := computeCacheKey("p", "europe-west1", "gemini-1.5-pro", c, s, nil) + otherProject, _ := computeCacheKey("q", "us-central1", "gemini-1.5-pro", c, s, nil) + assert.NotEqual(t, base, otherRegion) + assert.NotEqual(t, base, otherProject) +} + +// ----------------------------------------------------------------------- +// Integration tests: Resolve() against fake HTTP server +// ----------------------------------------------------------------------- + +func TestResolver_NoMarkers_ReturnsNil(t *testing.T) { + // Credentials must be set, or Resolve returns nil before it looks at the markers. + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, "", http.StatusOK) + r := resolverWithServer(fake.srv.URL) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{userMsg("hello")}, + } + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + assert.Nil(t, res, "no cache_control markers should return nil") + assert.Equal(t, 0, fake.creates()) +} + +// A resolver built without credentials is inert: even a marked request is not cached. +func TestResolver_NilAuth_IsInert(t *testing.T) { + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, "", http.StatusOK) + r := resolverWithServer(fake.srv.URL) + r.SetAuth(nil) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + assert.Nil(t, res) + assert.Equal(t, 0, fake.creates()) +} + +func TestResolver_CacheMiss_Creates(t *testing.T) { + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `","usageMetadata":{"totalTokenCount":512}}` + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) + r := resolverWithServer(fake.srv.URL) + + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, + } + + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new", res.CacheName) + assert.True(t, res.Created) + assert.Equal(t, 512, res.TokenCount) + // remainder = messages after the breakpoint (index 0 → only userMsg at index 1 remains) + assert.Equal(t, []openai.ChatCompletionMessageParamUnion{userMsg("Hello")}, res.Messages) + // The request path creates directly; it never lists. + assert.Equal(t, 0, fake.lists()) + assert.Equal(t, 1, fake.creates()) +} + +// A cache that exists in Google but not in the store is not discovered on the request +// path: the resolver creates rather than scanning cachedContents. Repairing that drift is +// the syncer's job. +func TestResolver_RequestPath_NeverLists(t *testing.T) { + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, + } + key := computeKeyForRequest(t, req) + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + + listBody, err := json.Marshal(map[string]interface{}{ + "cachedContents": []map[string]interface{}{ + { + "name": "projects/p/locations/us-central1/cachedContents/existing", + "displayName": key, + "model": "publishers/google/models/gemini-1.5-pro", + "expireTime": expireISO, + }, + }, + }) + require.NoError(t, err) + + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `"}` + fake := newFakeCacheServer(t, string(listBody), createResp, http.StatusOK) + r := resolverWithServer(fake.srv.URL) + + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new", res.CacheName) + assert.True(t, res.Created) + assert.Equal(t, 0, fake.lists(), "the request path must not list cachedContents") + assert.Equal(t, 1, fake.creates()) +} + +func TestResolver_StoreHit_SkipsGoogleAPICalls(t *testing.T) { + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, "", http.StatusOK) + store := newMemStore() + r := resolverWithServer(fake.srv.URL, store) + + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("Cached system.", ephemeralFields()), + userMsg("q"), + }, + } + key := computeKeyForRequest(t, req) + store.seed(key, "projects/p/locations/r/cachedContents/store-hit", time.Now().Add(10*time.Minute)) + + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/r/cachedContents/store-hit", res.CacheName) + assert.False(t, res.Created) + assert.Equal(t, 0, fake.lists(), "store hit must not call the Google API") + assert.Equal(t, 0, fake.creates()) +} + +func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + key := computeKeyForRequest(t, req) + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/refreshed","expireTime":"` + expireISO + `"}` + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) + store := newMemStore() + r := resolverWithServer(fake.srv.URL, store) + // Seed with an entry expiring inside the stale window → must be treated as a miss. + store.seed(key, "projects/p/locations/us-central1/cachedContents/stale", time.Now().Add(5*time.Second)) + + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/refreshed", res.CacheName) + assert.Equal(t, 1, fake.creates(), "a near-expiry store entry must be treated as a miss") + assert.Equal(t, 0, fake.lists()) + + // The fresh entry replaces the stale one in the store. + e, ok, err := store.Get(context.Background(), key) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/refreshed", e.Name) +} + +// A create response without an expireTime cannot be stored (its TTL would be negative), +// so every later request would create again. It must fail instead of passing silently. +func TestResolver_CreateExpiryMissing_Fails(t *testing.T) { + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, + `{"name":"projects/p/locations/us-central1/cachedContents/x"}`, http.StatusOK) + store := newMemStore() + r := resolverWithServer(fake.srv.URL, store) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + + res, err := r.Resolve(context.Background(), req) + require.Error(t, err) + assert.Nil(t, res) + assert.Contains(t, err.Error(), "no expireTime") +} + +// An unparseable expireTime fails JSON decoding of the create response. +func TestResolver_CreateExpiryUnparseable_Fails(t *testing.T) { + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, + `{"name":"projects/p/locations/us-central1/cachedContents/x","expireTime":"not-a-time"}`, http.StatusOK) + r := resolverWithServer(fake.srv.URL) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + + res, err := r.Resolve(context.Background(), req) + require.Error(t, err) + assert.Nil(t, res) +} + +func TestResolver_CreateFailure_ReturnsError(t *testing.T) { + fake := newFakeCacheServer(t, + `{"cachedContents":[]}`, + `{"error":{"message":"below minimum tokens"}}`, + http.StatusUnprocessableEntity, + ) + r := resolverWithServer(fake.srv.URL) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + _, err := r.Resolve(context.Background(), req) + require.Error(t, err) + assert.Contains(t, err.Error(), "422") +} + +func TestResolver_TTLDefault_SentToGoogle(t *testing.T) { + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + var capturedBody []byte + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"cachedContents":[]}`)) + return + } + buf := make([]byte, 8192) + n, _ := r.Body.Read(buf) + capturedBody = buf[:n] + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"name":"projects/p/locations/r/cachedContents/x","expireTime":"` + expireISO + `"}`)) + })) + defer srv.Close() + + r := resolverWithServer(srv.URL) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), // no TTL → default 300s + userMsg("q"), + }, + } + _, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + + var body gcp.CreateCachedContent + require.NoError(t, json.Unmarshal(capturedBody, &body)) + assert.Equal(t, "300s", body.TTL) + // The resolver sets a TTL, not an expiry. A zero-value expiry must not be sent, or + // Google would receive "0001-01-01T00:00:00Z" alongside the TTL. + assert.Nil(t, body.ExpireTime) + assert.NotContains(t, string(capturedBody), "expireTime") +} + +func TestResolver_TTLOverride_SentToGoogle(t *testing.T) { + expireISO := time.Now().Add(time.Hour).UTC().Format(time.RFC3339) + var capturedBody []byte + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"cachedContents":[]}`)) + return + } + buf := make([]byte, 8192) + n, _ := r.Body.Read(buf) + capturedBody = buf[:n] + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"name":"projects/p/locations/r/cachedContents/x","expireTime":"` + expireISO + `"}`)) + })) + defer srv.Close() + + r := resolverWithServer(srv.URL) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFieldsWithTTL("1h")), + userMsg("q"), + }, + } + _, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + + var body gcp.CreateCachedContent + require.NoError(t, json.Unmarshal(capturedBody, &body)) + assert.Equal(t, "3600s", body.TTL) +} + +// A replica keeps the cache it created, even if another replica created one for the same +// key. The list handler would return the other replica's cache on any call, so a +// reintroduced list (pre- or post-create) fails both the name and the list-count checks. +func TestResolver_DuplicateCreateRace_KeepsOwnCreate(t *testing.T) { + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + key := computeKeyForRequest(t, req) + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + + var listCalls atomic.Int64 + listBodyOther, _ := json.Marshal(map[string]interface{}{ + "cachedContents": []map[string]interface{}{ + { + "name": "projects/p/locations/us-central1/cachedContents/older-by-other", + "displayName": key, + "model": "publishers/google/models/gemini-1.5-pro", + "expireTime": expireISO, + }, + }, + }) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.Method { + case http.MethodGet: + listCalls.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(listBodyOther) + case http.MethodPost: + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"name":"projects/p/locations/us-central1/cachedContents/new-by-us","expireTime":"` + expireISO + `","usageMetadata":{"totalTokenCount":512}}`)) + } + })) + defer srv.Close() + + r := resolverWithServer(srv.URL) + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new-by-us", res.CacheName) + // The replica performed and paid for the create, so it must report it. + assert.True(t, res.Created) + assert.Equal(t, 512, res.TokenCount) + assert.Equal(t, int64(0), listCalls.Load(), "the request path must not list cachedContents") +} + +// The default TTL given to New is sent when the marker has none. +func TestResolver_ConfiguredDefaultTTL_SentToGoogle(t *testing.T) { + var body []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ = io.ReadAll(r.Body) + _, _ = fmt.Fprintf(w, `{"name":"c/x","expireTime":%q}`, time.Now().Add(10*time.Minute).UTC().Format(time.RFC3339)) + })) + t.Cleanup(srv.Close) + r := New(&http.Client{Transport: &redirectTransport{fakeHost: srv.URL[len("http://"):]}}, nil, syncAuth, 10*time.Minute, nil) + _, err := r.Resolve(context.Background(), &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{systemMsg("sys", ephemeralFields()), userMsg("q")}, + }) + require.NoError(t, err) + var got gcp.CreateCachedContent + require.NoError(t, json.Unmarshal(body, &got)) + assert.Equal(t, "600s", got.TTL) +} + +// Maps inside the key input (tool schema properties, function call args) must not make +// the key depend on map iteration order. +func TestComputeCacheKey_StableWithMaps(t *testing.T) { + props := map[string]*genai.Schema{} + args := map[string]any{} + for i := range 50 { + props[fmt.Sprintf("p%02d", i)] = &genai.Schema{Type: genai.TypeString} + args[fmt.Sprintf("a%02d", i)] = i + } + tools := []genai.Tool{{FunctionDeclarations: []*genai.FunctionDeclaration{{ + Name: "f", Parameters: &genai.Schema{Type: genai.TypeObject, Properties: props}, + }}}} + contents := []genai.Content{{Role: "model", Parts: []*genai.Part{{FunctionCall: &genai.FunctionCall{Name: "f", Args: args}}}}} + + want, err := computeCacheKey("p", "r", "m", contents, nil, tools) + require.NoError(t, err) + for range 100 { + got, err := computeCacheKey("p", "r", "m", contents, nil, tools) + require.NoError(t, err) + require.Equal(t, want, got) + } +} + +// A miss logs the miss and the create; the following request logs a hit. +func TestResolver_DebugLogs(t *testing.T) { + var buf bytes.Buffer + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug})) + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, + `{"name":"c/new","expireTime":"`+expireISO+`","usageMetadata":{"totalTokenCount":7}}`, http.StatusOK) + host := fake.srv.URL[len("http://"):] + r := New(&http.Client{Transport: &redirectTransport{fakeHost: host}}, newMemStore(), syncAuth, 0, logger) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{systemMsg("sys", ephemeralFields()), userMsg("q")}, + } + + _, err := r.Resolve(context.Background(), req) + require.NoError(t, err) + assert.Contains(t, buf.String(), "gcpcache: cache miss, creating") + assert.Contains(t, buf.String(), "gcpcache: cache created") + assert.Contains(t, buf.String(), "cache=c/new") + + buf.Reset() + _, err = r.Resolve(context.Background(), req) + require.NoError(t, err) + assert.Contains(t, buf.String(), "gcpcache: cache hit") +} diff --git a/internal/contextcache/gcpcache/redis_test.go b/internal/contextcache/gcpcache/redis_test.go new file mode 100644 index 0000000000..7630660994 --- /dev/null +++ b/internal/contextcache/gcpcache/redis_test.go @@ -0,0 +1,158 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package gcpcache + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/contextcache/redis" +) + +// ----------------------------------------------------------------------- +// Cross-replica behavior +// ----------------------------------------------------------------------- + +// Two resolvers stand in for two gateway replicas: separate processes, one shared Redis. +// Once one replica has published a cache name, the other reuses it without calling Google. +func TestResolver_CrossReplica_SecondResolveHitsStore(t *testing.T) { + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `","usageMetadata":{"totalTokenCount":512}}` + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) + + mr := miniredis.RunT(t) + newReplica := func() *Resolver { + store, err := redis.NewStore(mr.Addr()) + require.NoError(t, err) + return resolverWithServer(fake.srv.URL, store) + } + a, b := newReplica(), newReplica() + + resA, err := a.Resolve(context.Background(), crossReplicaRequest()) + require.NoError(t, err) + require.NotNil(t, resA) + assert.True(t, resA.Created) + assert.Equal(t, 512, resA.TokenCount) + + resB, err := b.Resolve(context.Background(), crossReplicaRequest()) + require.NoError(t, err) + require.NotNil(t, resB) + assert.Equal(t, resA.CacheName, resB.CacheName) + assert.False(t, resB.Created, "a store hit must not report a create") + assert.Equal(t, 0, resB.TokenCount) + + assert.Equal(t, 1, fake.creates(), "the second replica must reuse the published cache") + assert.Equal(t, 0, fake.lists()) +} + +// Replicas that miss the same cold prefix at the same moment each create a cache. This is +// the accepted cost of not coordinating: each replica keeps the cache it created, and each +// reports the create so the duplicate write is visible in cache-write token metrics. +func TestResolver_CrossReplica_ColdPrefixRaceMayDuplicate(t *testing.T) { + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + + // Hold every create until both replicas have issued one, so both miss the store. + const replicas = 2 + var creates atomic.Int64 + arrived := make(chan struct{}, replicas) + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("unexpected %s: the request path must only create", r.Method) + w.WriteHeader(http.StatusInternalServerError) + return + } + n := creates.Add(1) + arrived <- struct{}{} + <-release + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = fmt.Fprintf(w, `{"name":"projects/p/locations/us-central1/cachedContents/new-%d","expireTime":"%s","usageMetadata":{"totalTokenCount":512}}`, n, expireISO) + })) + t.Cleanup(srv.Close) + + mr := miniredis.RunT(t) + + var done sync.WaitGroup + results := make([]*contextcache.ResolveResult, replicas) + errs := make([]error, replicas) + for i := range replicas { + store, err := redis.NewStore(mr.Addr()) + require.NoError(t, err) + r := resolverWithServer(srv.URL, store) + done.Add(1) + go func() { + defer done.Done() + results[i], errs[i] = r.Resolve(context.Background(), crossReplicaRequest()) + }() + } + for range replicas { + <-arrived + } + close(release) + done.Wait() + + assert.Equal(t, int64(replicas), creates.Load()) + names := map[string]bool{} + for i, res := range results { + require.NoError(t, errs[i], "replica %d", i) + require.NotNil(t, res, "replica %d", i) + assert.True(t, res.Created, "replica %d performed a create and must report it", i) + assert.Equal(t, 512, res.TokenCount, "replica %d", i) + names[res.CacheName] = true + } + assert.Len(t, names, replicas, "each replica must keep the cache it created") +} + +func crossReplicaRequest() *openai.ChatCompletionRequest { + return &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, + } +} + +// The failure policy: an unreachable store degrades caching, it does not fail requests. +// The resolver falls back to Google and the request is served normally. +func TestResolver_StoreUnreachable_FailsOpen(t *testing.T) { + expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `","usageMetadata":{"totalTokenCount":512}}` + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) + + mr := miniredis.RunT(t) + store, err := redis.NewStore(mr.Addr()) + require.NoError(t, err) + mr.Close() // Every operation now fails to dial. + + r := resolverWithServer(fake.srv.URL, store) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, + } + + res, err := r.Resolve(context.Background(), req) + require.NoError(t, err, "a dead store must not fail the request") + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new", res.CacheName) + assert.Equal(t, 1, fake.creates(), "the resolver falls through to Google") +} diff --git a/internal/contextcache/gcpcache/syncer.go b/internal/contextcache/gcpcache/syncer.go new file mode 100644 index 0000000000..3a725f2243 --- /dev/null +++ b/internal/contextcache/gcpcache/syncer.go @@ -0,0 +1,185 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package gcpcache + +import ( + "context" + "encoding/hex" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "strconv" + "time" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/json" +) + +// This file supplies the GCP side of syncing: a contextcache.Source that walks +// cachedContents.list for one project and region. The loop, gate, and store writes live +// in contextcache.Syncer. + +const ( + // listPageSize is the documented maximum for cachedContents.list; larger values are + // coerced down by the API. Requesting it explicitly avoids an unspecified default. + listPageSize = 1000 + + // syncGatePrefix namespaces gate keys. Cache keys are 64 hex characters, so the + // prefix cannot collide with them. + syncGatePrefix = "gcpcache:sync:" +) + +// cachedContentItem is the subset of a cachedContents list item the syncer uses. +type cachedContentItem struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + ExpireTime string `json:"expireTime"` // RFC 3339 +} + +type listResponse struct { + CachedContents []cachedContentItem `json:"cachedContents"` + NextPageToken string `json:"nextPageToken"` +} + +// newSyncer returns a syncer. With a nil gcpAuth the source reads the resolver's current +// credentials each round, so SetAuth takes effect on the next round; a non-nil gcpAuth +// pins the credentials (used by tests). +func (r *Resolver) newSyncer(gcpAuth filterapi.GCPAuthHandler, interval time.Duration) *contextcache.Syncer { + return &contextcache.Syncer{ + Store: r.store, + Source: &listSource{r: r, auth: gcpAuth}, + Interval: interval, + Logger: r.logger, + } +} + +// listSource is a contextcache.Source over cachedContents.list. +type listSource struct { + r *Resolver + auth filterapi.GCPAuthHandler + // pages is the number of pages read by the last List call, for tests. + pages int +} + +// creds returns the pinned credentials, or the resolver's current ones. +func (s *listSource) creds() filterapi.GCPAuthHandler { + if s.auth != nil { + return s.auth + } + return s.r.GetAuth() +} + +// GateKey implements contextcache.Source. +func (s *listSource) GateKey() string { + auth := s.creds() + if auth == nil { + return syncGatePrefix + } + return syncGatePrefix + auth.GCPProject() + "/" + auth.GCPRegion() +} + +// List implements contextcache.Source. The API has no server-side filter, so the walk +// reads every entry in the project and region; only gateway-created ones are added. +func (s *listSource) List(ctx context.Context, add func(string, contextcache.Entry)) error { + s.pages = 0 + auth := s.creds() + if auth == nil { + return nil + } + region, project := auth.GCPRegion(), auth.GCPProject() + + token, err := auth.GCPTokenSource().Token() + if err != nil { + return fmt.Errorf("get GCP access token: %w", err) + } + base, err := url.Parse(fmt.Sprintf(gcpCachedContentsBasePath, region, project, region)) + if err != nil { + return fmt.Errorf("parse cachedContents URL: %w", err) + } + + pageToken := "" + for { + lr, err := s.r.listPage(ctx, base, token.AccessToken, pageToken) + if err != nil { + return err + } + s.pages++ + for _, item := range lr.CachedContents { + if key, e, ok := s.toEntry(item); ok { + add(key, e) + } + } + if lr.NextPageToken == "" { + return nil + } + pageToken = lr.NextPageToken + } +} + +// toEntry converts a list item. It reports false for items the gateway did not create +// (their displayName is not a cache key) and for items whose expiry does not parse. +func (s *listSource) toEntry(item cachedContentItem) (string, contextcache.Entry, bool) { + if !isCacheKey(item.DisplayName) || item.Name == "" { + return "", contextcache.Entry{}, false + } + expireTime, err := time.Parse(time.RFC3339, item.ExpireTime) + if err != nil { + s.r.logger.Warn("gcpcache: skipping cachedContents entry with unparseable expireTime", + slog.String("name", item.Name), slog.String("expireTime", item.ExpireTime)) + return "", contextcache.Entry{}, false + } + return item.DisplayName, contextcache.Entry{Name: item.Name, ExpireTime: expireTime}, true +} + +// listPage fetches one page of cachedContents. It sets the query on a copy of base, so +// the caller's URL is not altered. +func (r *Resolver) listPage(ctx context.Context, base *url.URL, accessToken, pageToken string) (listResponse, error) { + u := *base + q := url.Values{} + q.Set("pageSize", strconv.Itoa(listPageSize)) + if pageToken != "" { + q.Set("pageToken", pageToken) + } + u.RawQuery = q.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return listResponse{}, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + + resp, err := r.httpClient.Do(req) + if err != nil { + return listResponse{}, err + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return listResponse{}, err + } + if resp.StatusCode != http.StatusOK { + return listResponse{}, fmt.Errorf("list cachedContents returned HTTP %d: %s", resp.StatusCode, body) + } + + var lr listResponse + if err = json.Unmarshal(body, &lr); err != nil { + return listResponse{}, fmt.Errorf("failed to decode list response: %w", err) + } + return lr, nil +} + +// isCacheKey reports whether s has the shape computeCacheKey produces: a hex-encoded +// SHA-256 digest. +func isCacheKey(s string) bool { + if len(s) != 64 { + return false + } + _, err := hex.DecodeString(s) + return err == nil +} diff --git a/internal/contextcache/gcpcache/syncer_test.go b/internal/contextcache/gcpcache/syncer_test.go new file mode 100644 index 0000000000..c415675c71 --- /dev/null +++ b/internal/contextcache/gcpcache/syncer_test.go @@ -0,0 +1,311 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package gcpcache + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/contextcache/redis" +) + +var syncAuth = &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + +// newTestRedisStore starts an in-process Redis and returns a store pointed at it. +func newTestRedisStore(t *testing.T) (contextcache.Store, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + s, err := redis.NewStore(mr.Addr()) + require.NoError(t, err) + return s, mr +} + +// syncOnce runs one round for auth against r's store and returns its stats and the +// number of list pages read. +func syncOnce(t *testing.T, r *Resolver, auth *fakeGCPAuth) (contextcache.Stats, int, error) { + t.Helper() + rc := r.newSyncer(auth, time.Minute) + stats, err := rc.RunOnce(context.Background()) + return stats, rc.Source.(*listSource).pages, err +} + +func testKey(n int) string { return fmt.Sprintf("%064x", n) } + +func itemJSON(key, name string, expire time.Time) string { + return fmt.Sprintf(`{"name":%q,"displayName":%q,"expireTime":%q}`, name, key, expire.UTC().Format(time.RFC3339)) +} + +// newListServer serves pages keyed by pageToken ("" is the first page). It fails the test +// on any non-GET, and records every query it saw. +func newListServer(t *testing.T, pages map[string]string) (*httptest.Server, *atomic.Int64, *[]string) { + t.Helper() + var calls atomic.Int64 + var queries []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + t.Errorf("syncer issued %s; it must only list", r.Method) + } + calls.Add(1) + queries = append(queries, r.URL.RawQuery) + body, ok := pages[r.URL.Query().Get("pageToken")] + if !ok { + w.WriteHeader(http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(srv.Close) + return srv, &calls, &queries +} + +func TestSync_MultiPageWalkWritesEveryEntry(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + srv, calls, queries := newListServer(t, map[string]string{ + "": `{"cachedContents":[` + itemJSON(testKey(1), "c/1", exp) + `],"nextPageToken":"p2"}`, + "p2": `{"cachedContents":[` + itemJSON(testKey(2), "c/2", exp) + `]}`, + }) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + stats, pages, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + assert.Equal(t, 2, pages) + assert.Equal(t, 2, stats.Written) + assert.Equal(t, int64(2), calls.Load()) + for _, q := range *queries { + assert.Contains(t, q, "pageSize=1000") + } + for n, name := range map[int]string{1: "c/1", 2: "c/2"} { + e, ok, err := store.Get(context.Background(), testKey(n)) + require.NoError(t, err) + require.True(t, ok, "key %d", n) + assert.Equal(t, name, e.Name) + } +} + +func TestSync_DoesNotOverwriteExistingEntry(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + srv, _, _ := newListServer(t, map[string]string{ + "": `{"cachedContents":[` + itemJSON(testKey(1), "c/from-list", exp) + `]}`, + }) + store, _ := newTestRedisStore(t) + require.NoError(t, store.Set(context.Background(), testKey(1), contextcache.Entry{Name: "c/published", ExpireTime: exp}, time.Minute)) + r := resolverWithServer(srv.URL, store) + + stats, _, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + assert.Equal(t, 0, stats.Written) + e, _, _ := store.Get(context.Background(), testKey(1)) + assert.Equal(t, "c/published", e.Name, "syncer must not clobber a request's write") +} + +func TestSync_GateSkipsSecondReplica(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, mr := newTestRedisStore(t) + second, err := redis.NewStore(mr.Addr()) + require.NoError(t, err) + a := resolverWithServer(srv.URL, store) + b := resolverWithServer(srv.URL, second) + + statsA, _, err := syncOnce(t, a, syncAuth) + require.NoError(t, err) + assert.False(t, statsA.Gated) + statsB, _, err := syncOnce(t, b, syncAuth) + require.NoError(t, err) + assert.True(t, statsB.Gated) + assert.Equal(t, int64(1), calls.Load()) + + // Once the interval passes, a replica may sync again. + mr.FastForward(time.Minute + time.Second) + statsB, _, err = syncOnce(t, b, syncAuth) + require.NoError(t, err) + assert.False(t, statsB.Gated) +} + +func TestSync_GateIsPerProjectRegion(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + other := &fakeGCPAuth{token: "tok", region: "europe-west1", project: "p"} + + _, _, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + stats, _, err := syncOnce(t, r, other) + require.NoError(t, err) + assert.False(t, stats.Gated, "one region's round must not suppress another's") + assert.Equal(t, int64(2), calls.Load()) +} + +func TestSync_SkipsForeignExpiredAndUnparseable(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + items := []string{ + itemJSON("my-own-cache", "c/foreign", exp), // not a cache key + itemJSON(testKey(1), "c/near-expiry", time.Now().Add(time.Second)), // within staleThreshold + fmt.Sprintf(`{"name":"c/bad","displayName":%q,"expireTime":"not-a-time"}`, testKey(2)), + itemJSON(testKey(3), "c/good", exp), + } + srv, _, _ := newListServer(t, map[string]string{"": `{"cachedContents":[` + strings.Join(items, ",") + `]}`}) + store, mr := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + stats, _, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + // Foreign and unparseable items are dropped by the source; the near-expiry one is + // seen but not written. + assert.Equal(t, 2, stats.Seen) + assert.Equal(t, 1, stats.Written) + assert.Equal(t, 1, stats.Skipped) + assert.False(t, mr.Exists("my-own-cache")) + assert.False(t, mr.Exists(testKey(1))) + assert.False(t, mr.Exists(testKey(2))) + assert.True(t, mr.Exists(testKey(3))) +} + +func TestSync_EntryTTLMatchesExpiry(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + srv, _, _ := newListServer(t, map[string]string{"": `{"cachedContents":[` + itemJSON(testKey(1), "c/1", exp) + `]}`}) + store, mr := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + _, _, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + assert.InDelta(t, (5 * time.Minute).Seconds(), mr.TTL(testKey(1)).Seconds(), 5) +} + +func TestSync_ListFailureReturnsError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + _, _, err := syncOnce(t, r, syncAuth) + require.ErrorContains(t, err, "HTTP 500") +} + +// A no-op store never wins the gate, so its rounds never list the provider. +func TestSync_NoopStoreIsGatedAndNeverLists(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + r := resolverWithServer(srv.URL) + + stats, _, err := syncOnce(t, r, syncAuth) + require.NoError(t, err) + assert.True(t, stats.Gated) + assert.Equal(t, int64(0), calls.Load()) +} + +func TestIsCacheKey(t *testing.T) { + assert.True(t, isCacheKey(testKey(7))) + assert.False(t, isCacheKey("")) + assert.False(t, isCacheKey(strings.Repeat("z", 64))) + assert.False(t, isCacheKey(testKey(7)[:63])) +} + +// ----------------------------------------------------------------------- +// Lifecycle +// ----------------------------------------------------------------------- + +func TestResolver_StartRunsSyncerAndCloseStopsIt(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + srv, calls, _ := newListServer(t, map[string]string{ + "": `{"cachedContents":[` + itemJSON(testKey(1), "c/1", exp) + `]}`, + }) + store, mr := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + r.Start() + require.Eventually(t, func() bool { return mr.Exists(testKey(1)) }, 2*time.Second, 10*time.Millisecond, + "Start must run a sync round") + + require.NoError(t, r.Close()) + select { + case <-r.done: + default: + t.Fatal("Close must wait for the syncer goroutine to exit") + } + assert.Equal(t, int64(1), calls.Load()) +} + +func TestResolver_StartIsIdempotent(t *testing.T) { + srv, _, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + r.Start() + first := r.done + r.Start() + assert.Equal(t, first, r.done, "a second Start must not launch another syncer") + require.NoError(t, r.Close()) +} + +func TestResolver_CloseIsIdempotentAndBlocksLaterStart(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + require.NoError(t, r.Close(), "Close before Start must be safe") + require.NoError(t, r.Close()) + r.Start() + assert.Nil(t, r.cancel, "Start after Close must do nothing") + assert.Equal(t, int64(0), calls.Load()) +} + +func TestResolver_StartWithNoopStoreClosesCleanly(t *testing.T) { + r := New(nil, nil, syncAuth, 0, nil) + r.Start() + require.NoError(t, r.Close()) + select { + case <-r.done: + default: + t.Fatal("Close must wait for the syncer goroutine to exit") + } +} + +// The syncer reads the resolver's current credentials each round, so SetAuth after a +// config reload moves the next round to the new project and region. +func TestSync_FollowsSetAuth(t *testing.T) { + srv, _, queries := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + rc := r.newSyncer(nil, time.Minute) + + assert.Equal(t, syncGatePrefix+"p/us-central1", rc.Source.GateKey()) + r.SetAuth(&fakeGCPAuth{token: "tok2", region: "europe-west1", project: "q"}) + assert.Equal(t, syncGatePrefix+"q/europe-west1", rc.Source.GateKey()) + + _, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + require.Len(t, *queries, 1) +} + +// Resolve uses the credentials set most recently, so a rotated token reaches GCP. +func TestResolver_SetAuthRotatesToken(t *testing.T) { + var gotAuth atomic.Value + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth.Store(r.Header.Get("Authorization")) + _, _ = fmt.Fprintf(w, `{"name":"c/x","expireTime":%q}`, time.Now().Add(5*time.Minute).UTC().Format(time.RFC3339)) + })) + t.Cleanup(srv.Close) + r := resolverWithServer(srv.URL) + r.SetAuth(&fakeGCPAuth{token: "rotated", region: "us-central1", project: "p"}) + + _, err := r.Resolve(context.Background(), crossReplicaRequest()) + require.NoError(t, err) + assert.Equal(t, "Bearer rotated", gotAuth.Load()) +} diff --git a/internal/contextcache/redis/store.go b/internal/contextcache/redis/store.go new file mode 100644 index 0000000000..6443b1138d --- /dev/null +++ b/internal/contextcache/redis/store.go @@ -0,0 +1,115 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package redis + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" +) + +// redisStore is a contextcache.Store backed by Redis, shared across gateway replicas. +type redisStore struct { + client goredis.UniversalClient +} + +// NewStore returns a contextcache.Store backed by the Redis instance at url, which accepts +// either a bare "host:port" or a full "redis://" URL. +// +// No connection is established here; go-redis dials lazily. A Redis that is unreachable +// therefore surfaces as an error from Get/Set, which resolvers treat as a miss. +func NewStore(url string) (contextcache.Store, error) { + opts, err := parseRedisURL(url) + if err != nil { + return nil, err + } + return &redisStore{client: goredis.NewClient(opts)}, nil +} + +// parseRedisURL accepts both a scheme-qualified URL and a bare host:port. +func parseRedisURL(url string) (*goredis.Options, error) { + if url == "" { + return nil, errors.New("contextcache: redis url is empty") + } + if strings.Contains(url, "://") { + opts, err := goredis.ParseURL(url) + if err != nil { + return nil, fmt.Errorf("contextcache: invalid redis url: %w", err) + } + return opts, nil + } + return &goredis.Options{Addr: url}, nil +} + +// Get implements Store. +func (s *redisStore) Get(ctx context.Context, key string) (contextcache.Entry, bool, error) { + v, err := s.client.Get(ctx, key).Result() + if err != nil { + if errors.Is(err, goredis.Nil) { + return contextcache.Entry{}, false, nil + } + return contextcache.Entry{}, false, fmt.Errorf("contextcache: redis get: %w", err) + } + e, err := decodeEntry(v) + if err != nil { + // A malformed value is treated as a miss rather than an error: it is not worth + // failing a resolution over, and the next write will overwrite it. + return contextcache.Entry{}, false, nil + } + return e, true, nil +} + +// Set implements Store. +func (s *redisStore) Set(ctx context.Context, key string, e contextcache.Entry, ttl time.Duration) error { + if err := s.client.Set(ctx, key, encodeEntry(e), ttl).Err(); err != nil { + return fmt.Errorf("contextcache: redis set: %w", err) + } + return nil +} + +// SetNX implements Store. +func (s *redisStore) SetNX(ctx context.Context, key string, e contextcache.Entry, ttl time.Duration) (bool, error) { + wrote, err := s.client.SetNX(ctx, key, encodeEntry(e), ttl).Result() + if err != nil { + return false, fmt.Errorf("contextcache: redis setnx: %w", err) + } + return wrote, nil +} + +// AcquireGate implements Store. +func (s *redisStore) AcquireGate(ctx context.Context, key string, ttl time.Duration) (bool, error) { + won, err := s.client.SetNX(ctx, key, "1", ttl).Result() + if err != nil { + return false, fmt.Errorf("contextcache: redis gate: %w", err) + } + return won, nil +} + +var _ contextcache.Store = (*redisStore)(nil) + +// encodeEntry serializes an entry as "|". The cache name is +// last because it is the only field that may itself contain the separator. +func encodeEntry(e contextcache.Entry) string { + return e.ExpireTime.UTC().Format(time.RFC3339) + "|" + e.Name +} + +func decodeEntry(s string) (contextcache.Entry, error) { + expiry, name, ok := strings.Cut(s, "|") + if !ok || name == "" { + return contextcache.Entry{}, fmt.Errorf("contextcache: malformed cache entry %q", s) + } + t, err := time.Parse(time.RFC3339, expiry) + if err != nil { + return contextcache.Entry{}, fmt.Errorf("contextcache: malformed cache entry expiry %q: %w", expiry, err) + } + return contextcache.Entry{Name: name, ExpireTime: t}, nil +} diff --git a/internal/contextcache/redis/store_test.go b/internal/contextcache/redis/store_test.go new file mode 100644 index 0000000000..7b1569e596 --- /dev/null +++ b/internal/contextcache/redis/store_test.go @@ -0,0 +1,97 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package redis + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" +) + +// newTestRedisStore starts an in-process Redis and returns a store pointed at it. +func newTestRedisStore(t *testing.T) (*redisStore, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + s, err := NewStore(mr.Addr()) + require.NoError(t, err) + return s.(*redisStore), mr +} + +func TestParseRedisURL(t *testing.T) { + t.Run("bare host:port", func(t *testing.T) { + opts, err := parseRedisURL("localhost:6379") + require.NoError(t, err) + assert.Equal(t, "localhost:6379", opts.Addr) + }) + t.Run("scheme-qualified", func(t *testing.T) { + opts, err := parseRedisURL("redis://user:pw@localhost:6380/2") + require.NoError(t, err) + assert.Equal(t, "localhost:6380", opts.Addr) + assert.Equal(t, 2, opts.DB) + }) + t.Run("empty", func(t *testing.T) { + _, err := parseRedisURL("") + require.Error(t, err) + }) + t.Run("malformed", func(t *testing.T) { + _, err := parseRedisURL("http://localhost:6379") + require.Error(t, err) + }) +} + +func TestRedisStore_SetGetRoundTrip(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + expire := time.Now().Add(10 * time.Minute).UTC().Truncate(time.Second) + + require.NoError(t, s.Set(ctx, "k", contextcache.Entry{Name: "projects/p/locations/r/cachedContents/x", ExpireTime: expire}, time.Minute)) + + got, ok, err := s.Get(ctx, "k") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "projects/p/locations/r/cachedContents/x", got.Name) + assert.True(t, expire.Equal(got.ExpireTime), "want %s got %s", expire, got.ExpireTime) +} + +func TestRedisStore_GetMiss(t *testing.T) { + s, _ := newTestRedisStore(t) + _, ok, err := s.Get(context.Background(), "absent") + require.NoError(t, err) + assert.False(t, ok) +} + +// A value that does not decode is treated as a miss rather than an error: it is not worth +// failing a resolution over, and the next write overwrites it. +func TestRedisStore_MalformedValueIsAMiss(t *testing.T) { + s, mr := newTestRedisStore(t) + require.NoError(t, mr.Set("k", "not-an-entry")) + + _, ok, err := s.Get(context.Background(), "k") + require.NoError(t, err) + assert.False(t, ok) +} + +func TestEncodeDecodeEntry(t *testing.T) { + expire := time.Now().Add(time.Hour).UTC().Truncate(time.Second) + // The cache name is encoded last precisely because it may contain the separator. + e := contextcache.Entry{Name: "projects/p|weird/cachedContents/x", ExpireTime: expire} + + got, err := decodeEntry(encodeEntry(e)) + require.NoError(t, err) + assert.Equal(t, e.Name, got.Name) + assert.True(t, expire.Equal(got.ExpireTime)) + + for _, bad := range []string{"", "no-separator", "not-a-time|name", "2020-01-01T00:00:00Z|"} { + _, err := decodeEntry(bad) + assert.Error(t, err, "input %q", bad) + } +} diff --git a/internal/contextcache/resolver.go b/internal/contextcache/resolver.go new file mode 100644 index 0000000000..bc9f506018 --- /dev/null +++ b/internal/contextcache/resolver.go @@ -0,0 +1,43 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package contextcache + +import ( + "context" + "time" + + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" +) + +// ResolveResult holds the outcome of a successful cache resolution. +type ResolveResult struct { + // CacheName is the provider's name for the resolved or created cache. For GCP Vertex + // AI it is "projects/{project}/locations/{location}/cachedContents/{cache_id}". + CacheName string + // Messages is the non-cached remainder of the conversation (messages after the breakpoint). + Messages []openai.ChatCompletionMessageParamUnion + // Created is true when this call created a new cache entry (cache-write cost applies). + Created bool + // TokenCount is the number of tokens stored in the cache, from the provider's create + // response. Only populated when Created is true. + TokenCount int + // ExpireTime is the cache expiration time reported by the provider. + ExpireTime time.Time +} + +// CacheResolver resolves or creates a provider cache for requests that carry +// Anthropic-style cache_control markers. It is all the request path needs; lifecycle +// methods live on CacheSyncer so request code cannot stop syncing for a shared resolver. +// +// Providers bind their credentials when the implementation is constructed. +type CacheResolver interface { + // Resolve inspects openAIReq for cache_control markers, resolves or creates the + // corresponding provider cache, and returns the result. It returns (nil, nil) when + // the request has no markers. + // The caller is responsible for injecting ResolveResult.CacheName into the provider + // request and replacing the request messages with ResolveResult.Messages. + Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest) (*ResolveResult, error) +} diff --git a/internal/contextcache/store.go b/internal/contextcache/store.go new file mode 100644 index 0000000000..6e4b088a6c --- /dev/null +++ b/internal/contextcache/store.go @@ -0,0 +1,81 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +// Package contextcache holds the provider-neutral parts of context caching: the shared +// store that maps a deterministic cache key to a provider cache name, and the background +// syncer that repairs drift between that store and the provider. +// +// Provider packages (such as gcpcache) own everything that talks to a provider API: +// computing cache keys, creating caches, and listing them for the syncer. +// +// The shared store deduplicates callers separated in time: once any replica publishes a +// cache name, every replica reuses it. Replicas that miss the same cold prefix at the +// same moment each create their own cache; this is tolerated, and each create is billed +// and reported. +package contextcache + +import ( + "context" + "time" +) + +// Entry is a resolved provider cache: its name and when it expires. +type Entry struct { + // Name identifies the cache to the provider. For GCP Vertex AI it is the full + // cachedContents resource name. + Name string + // ExpireTime is when the provider cache expires. + ExpireTime time.Time +} + +// Store maps a deterministic cache key to the provider cache it resolved to. +// +// The store is shared across gateway replicas, so a cache name published by one +// replica is reused by all of them. Replicas that miss the +// same cold prefix at the same moment may each create a provider cache, and resolvers +// tolerate that. Implementations must be safe for concurrent use. +// +// Store errors are never fatal to a request: resolvers treat them as a cache miss and +// proceed to the provider. +type Store interface { + // Get returns the entry for key. The bool reports whether a usable entry was + // found; a miss returns (Entry{}, false, nil). + Get(ctx context.Context, key string) (Entry, bool, error) + // Set records e under key, expiring it after ttl. + Set(ctx context.Context, key string, e Entry, ttl time.Duration) error + // SetNX records e under key only if key is absent. It reports whether it wrote. + SetNX(ctx context.Context, key string, e Entry, ttl time.Duration) (bool, error) + // AcquireGate claims key for ttl. It reports whether the caller won. + AcquireGate(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +// NoopStore is the Store used when no shared store is configured. Every lookup misses +// and every write is discarded, so context caching becomes inert: markers are still +// parsed, but nothing is resolved or created. +// +// This is what makes "store unconfigured" behave identically to "store unreachable" +// without a nil check at each call site. +// +// AcquireGate never wins, so a Syncer over a NoopStore reports every round as gated +// and never lists the provider. +type NoopStore struct{} + +// Get implements Store. +func (NoopStore) Get(context.Context, string) (Entry, bool, error) { return Entry{}, false, nil } + +// Set implements Store. +func (NoopStore) Set(context.Context, string, Entry, time.Duration) error { return nil } + +// SetNX implements Store. +func (NoopStore) SetNX(context.Context, string, Entry, time.Duration) (bool, error) { + return false, nil +} + +// AcquireGate implements Store. +func (NoopStore) AcquireGate(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +var _ Store = NoopStore{} diff --git a/internal/contextcache/syncer.go b/internal/contextcache/syncer.go new file mode 100644 index 0000000000..25343d09a9 --- /dev/null +++ b/internal/contextcache/syncer.go @@ -0,0 +1,147 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package contextcache + +import ( + "context" + "fmt" + "io" + "log/slog" + "time" +) + +// CacheSyncer runs a provider's background sync. Providers bind their credentials when +// the implementation is constructed, so Start takes no arguments. +type CacheSyncer interface { + // Start begins background syncing. It runs under a context the implementation owns, + // so it outlives the call. It does nothing if already started or if closed. + Start() + // Close stops background syncing and waits for it to exit. It is idempotent. + io.Closer +} + +// StaleThreshold is how close to its expiry an entry must be before it is treated as a +// miss, so there is time to re-resolve before the provider cache disappears underneath +// an in-flight request. The syncer does not write entries this close to expiry. +const StaleThreshold = 10 * time.Second + +// Source lists the caches a provider holds for one scope, such as a GCP project and +// region. It is the only provider-specific part of syncing. +type Source interface { + // GateKey is the store key used to rate-limit rounds for this scope across the + // fleet. It must be distinct per scope, and must not collide with cache keys. + GateKey() string + // List calls add once per cache the gateway created in this scope, with the cache + // key it was created under. It should skip caches it cannot key or date. An error + // ends the round; entries already added stay written. + List(ctx context.Context, add func(key string, e Entry)) error +} + +// Stats summarizes one sync round. +type Stats struct { + // Gated is true when another replica held the gate and nothing was listed. + Gated bool + // Seen is the number of entries the source added. + Seen int + // Written is the number of entries written to the store. + Written int + // Skipped is the number of entries not written: already present, near expiry, or + // failed to write. + Skipped int +} + +// Syncer repairs drift between a Store and a provider: caches that exist at the +// provider but that the store does not know about, after a store flush or eviction, a +// dropped write, or a cache created before the store was configured. +// +// Each round: +// 1. Acquires the source's gate with SET NX and the interval as TTL, so across the fleet +// at most one replica lists a scope per interval. A replica that loses skips the round. +// 2. Lists the source. +// 3. Writes each entry with SET NX and a TTL pinned to its expiry. SET NX means it only +// fills gaps and never overwrites a name that a request just published. +type Syncer struct { + Store Store + Source Source + Interval time.Duration + Logger *slog.Logger +} + +// Run runs rounds until ctx is done. It never returns an error: a failed round is logged +// and the next one retries. +func (rc *Syncer) Run(ctx context.Context) { + ticker := time.NewTicker(rc.Interval) + defer ticker.Stop() + for { + rc.round(ctx) + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} + +func (rc *Syncer) round(ctx context.Context) { + stats, err := rc.RunOnce(ctx) + if err != nil { + if ctx.Err() == nil { + rc.logger().Warn("contextcache: sync round failed", + slog.String("gate", rc.Source.GateKey()), slog.String("error", err.Error())) + } + return + } + if !stats.Gated { + rc.logger().Debug("contextcache: sync round complete", + slog.String("gate", rc.Source.GateKey()), slog.Int("seen", stats.Seen), + slog.Int("written", stats.Written), slog.Int("skipped", stats.Skipped)) + } +} + +// RunOnce performs one round. An error ends the round early; entries written before it +// stay written. +func (rc *Syncer) RunOnce(ctx context.Context) (Stats, error) { + var stats Stats + won, err := rc.Store.AcquireGate(ctx, rc.Source.GateKey(), rc.Interval) + if err != nil { + return stats, fmt.Errorf("acquire sync gate: %w", err) + } + if !won { + stats.Gated = true + return stats, nil + } + err = rc.Source.List(ctx, func(key string, e Entry) { + stats.Seen++ + if rc.write(ctx, key, e) { + stats.Written++ + } else { + stats.Skipped++ + } + }) + return stats, err +} + +func (rc *Syncer) write(ctx context.Context, key string, e Entry) bool { + // Resolvers treat entries this close to expiry as a miss, so writing one is pointless. + ttl := time.Until(e.ExpireTime) + if ttl < StaleThreshold { + return false + } + wrote, err := rc.Store.SetNX(ctx, key, e, ttl) + if err != nil { + rc.logger().Warn("contextcache: sync store write failed", + slog.String("name", e.Name), slog.String("error", err.Error())) + return false + } + return wrote +} + +func (rc *Syncer) logger() *slog.Logger { + if rc.Logger == nil { + return slog.New(slog.DiscardHandler) + } + return rc.Logger +} diff --git a/internal/contextcache/syncer_test.go b/internal/contextcache/syncer_test.go new file mode 100644 index 0000000000..ff7d34c667 --- /dev/null +++ b/internal/contextcache/syncer_test.go @@ -0,0 +1,186 @@ +// Copyright Envoy AI Gateway Authors +// SPDX-License-Identifier: Apache-2.0 +// The full text of the Apache license is available in the LICENSE file at +// the root of the repo. + +package contextcache + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeSource adds a fixed set of entries, then returns err. +type fakeSource struct { + gate string + entries map[string]Entry + err error + calls int +} + +func (s *fakeSource) GateKey() string { return s.gate } + +func (s *fakeSource) List(_ context.Context, add func(string, Entry)) error { + s.calls++ + for k, e := range s.entries { + add(k, e) + } + return s.err +} + +// memStore is an in-memory Store for syncer tests. TTLs are ignored except that a gate, +// once taken, is never released, which is enough for a single test. +type memStore struct { + mu sync.Mutex + entries map[string]Entry + gates map[string]bool +} + +func newMemStore() *memStore { + return &memStore{entries: map[string]Entry{}, gates: map[string]bool{}} +} + +func (m *memStore) Get(_ context.Context, key string) (Entry, bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + e, ok := m.entries[key] + return e, ok, nil +} + +func (m *memStore) Set(_ context.Context, key string, e Entry, _ time.Duration) error { + m.mu.Lock() + defer m.mu.Unlock() + m.entries[key] = e + return nil +} + +func (m *memStore) SetNX(_ context.Context, key string, e Entry, _ time.Duration) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.entries[key]; ok { + return false, nil + } + m.entries[key] = e + return true, nil +} + +func (m *memStore) AcquireGate(_ context.Context, key string, _ time.Duration) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.gates[key] { + return false, nil + } + m.gates[key] = true + return true, nil +} + +func newTestSyncer(_ *testing.T, src *fakeSource) (*Syncer, Store) { + store := newMemStore() + return &Syncer{Store: store, Source: src, Interval: time.Minute}, store +} + +func TestSyncer_WritesEntries(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestSyncer(t, &fakeSource{gate: "g", entries: map[string]Entry{ + "a": {Name: "c/a", ExpireTime: exp}, + "b": {Name: "c/b", ExpireTime: exp}, + }}) + + stats, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + assert.Equal(t, Stats{Seen: 2, Written: 2}, stats) + e, ok, err := store.Get(context.Background(), "a") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "c/a", e.Name) +} + +func TestSyncer_DoesNotOverwrite(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestSyncer(t, &fakeSource{gate: "g", entries: map[string]Entry{ + "a": {Name: "c/from-list", ExpireTime: exp}, + }}) + require.NoError(t, store.Set(context.Background(), "a", Entry{Name: "c/published", ExpireTime: exp}, time.Minute)) + + stats, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + assert.Equal(t, Stats{Seen: 1, Skipped: 1}, stats) + e, _, _ := store.Get(context.Background(), "a") + assert.Equal(t, "c/published", e.Name) +} + +func TestSyncer_SkipsNearExpiry(t *testing.T) { + rc, store := newTestSyncer(t, &fakeSource{gate: "g", entries: map[string]Entry{ + "a": {Name: "c/a", ExpireTime: time.Now().Add(time.Second)}, + }}) + + stats, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + assert.Equal(t, Stats{Seen: 1, Skipped: 1}, stats) + _, ok, _ := store.Get(context.Background(), "a") + assert.False(t, ok) +} + +func TestSyncer_GateSkipsSecondRound(t *testing.T) { + src := &fakeSource{gate: "g"} + rc, _ := newTestSyncer(t, src) + + _, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + stats, err := rc.RunOnce(context.Background()) + require.NoError(t, err) + assert.True(t, stats.Gated) + assert.Equal(t, 1, src.calls, "a gated round must not list") +} + +func TestSyncer_ListErrorKeepsEarlierWrites(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestSyncer(t, &fakeSource{ + gate: "g", err: errors.New("boom"), + entries: map[string]Entry{"a": {Name: "c/a", ExpireTime: exp}}, + }) + + _, err := rc.RunOnce(context.Background()) + require.ErrorContains(t, err, "boom") + _, ok, _ := store.Get(context.Background(), "a") + assert.True(t, ok) +} + +func TestSyncer_RunStopsOnCancel(t *testing.T) { + rc, _ := newTestSyncer(t, &fakeSource{gate: "g", err: errors.New("boom")}) + rc.Interval = time.Hour + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + rc.Run(ctx) // A failing round must not stop or panic the loop. + close(done) + }() + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Run must return when its context is cancelled") + } +} + +func TestNoopStore_SyncerMethods(t *testing.T) { + var s NoopStore + wrote, err := s.SetNX(context.Background(), "k", Entry{Name: "n", ExpireTime: time.Now().Add(time.Hour)}, time.Minute) + require.NoError(t, err) + assert.False(t, wrote) + won, err := s.AcquireGate(context.Background(), "g", time.Minute) + require.NoError(t, err) + assert.False(t, won) + + src := &fakeSource{gate: "g"} + stats, err := (&Syncer{Store: s, Source: src, Interval: time.Minute}).RunOnce(context.Background()) + require.NoError(t, err) + assert.True(t, stats.Gated) + assert.Equal(t, 0, src.calls) +} diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 13f9b94156..19cdf0673d 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -9,6 +9,7 @@ import ( "context" "fmt" "strconv" + "time" egv1a1 "github.com/envoyproxy/gateway/api/v1alpha1" "github.com/go-logr/logr" @@ -70,6 +71,11 @@ type Options struct { ExtProcLogFormat string // ExtProcEnableRedaction enables redaction of sensitive information in debug logs for the external processor. ExtProcEnableRedaction bool + // ExtProcContextCacheStoreURL is the shared store for GCP Vertex AI context caching in the + // external processor. Empty disables context caching. + ExtProcContextCacheStoreURL string + // ExtProcContextCacheDefaultTTL is the cache TTL used when a cache_control marker gives none. + ExtProcContextCacheDefaultTTL time.Duration // ExtProcImage is the image for the external processor set on Deployment. ExtProcImage string // ExtProcImagePullPolicy is the image pull policy for the external processor set on Deployment. diff --git a/internal/controller/extproc_builder.go b/internal/controller/extproc_builder.go index f647ead190..4de4f45870 100644 --- a/internal/controller/extproc_builder.go +++ b/internal/controller/extproc_builder.go @@ -13,6 +13,7 @@ import ( "path/filepath" "strconv" "strings" + "time" egv1a1 "github.com/envoyproxy/gateway/api/v1alpha1" "github.com/go-logr/logr" @@ -60,6 +61,8 @@ func newExtProcBuilder(options *Options, extProcAsSideCar bool, logger logr.Logg logLevel: options.ExtProcLogLevel, logFormat: options.ExtProcLogFormat, enableRedaction: options.ExtProcEnableRedaction, + contextCacheStoreURL: options.ExtProcContextCacheStoreURL, + contextCacheDefaultTTL: options.ExtProcContextCacheDefaultTTL, udsPath: options.UDSPath, requestHeaderAttributes: options.RequestHeaderAttributes, spanRequestHeaderAttributes: options.TracingRequestHeaderAttributes, @@ -88,6 +91,9 @@ type extProcBuilder struct { enableRedaction bool udsPath string + contextCacheStoreURL string + contextCacheDefaultTTL time.Duration + requestHeaderAttributes *string spanRequestHeaderAttributes *string metricsRequestHeaderAttributes *string @@ -210,8 +216,12 @@ type extProcConfigDigest struct { LogLevel string // LogFormat is omitted when it is the default, so deployments that never asked for JSON keep // the same hash as before this field existed and are not rolled on upgrade. - LogFormat string `json:",omitempty"` - EnableRedaction bool + LogFormat string `json:",omitempty"` + EnableRedaction bool + // ContextCache* are omitted when caching is off, so deployments without it keep the + // same hash and are not rolled on upgrade. + ContextCacheStoreURL string `json:",omitempty"` + ContextCacheDefaultTTL time.Duration `json:",omitempty"` UDSPath string RequestHeaderAttributes *string SpanRequestHeaderAttributes *string @@ -239,6 +249,8 @@ func (b *extProcBuilder) extProcContainerHash(input extProcContainerInput) strin LogLevel: b.logLevel, LogFormat: b.nonDefaultLogFormat(), EnableRedaction: b.enableRedaction, + ContextCacheStoreURL: b.contextCacheStoreURL, + ContextCacheDefaultTTL: b.contextCacheTTLArg(), UDSPath: b.udsPath, RequestHeaderAttributes: b.requestHeaderAttributes, SpanRequestHeaderAttributes: b.spanRequestHeaderAttributes, @@ -330,9 +342,24 @@ func (b *extProcBuilder) buildExtProcBaseArgs(needMCP bool) []string { if b.enableRedaction { args = append(args, "-enableRedaction") } + if b.contextCacheStoreURL != "" { + args = append(args, "-contextCacheStoreURL", b.contextCacheStoreURL) + if ttl := b.contextCacheTTLArg(); ttl != 0 { + args = append(args, "-contextCacheDefaultTTL", ttl.String()) + } + } return args } +// contextCacheTTLArg returns the TTL to pass to extproc, or 0 when caching is off or the +// TTL is unset, in which case extproc uses its own default. +func (b *extProcBuilder) contextCacheTTLArg() time.Duration { + if b.contextCacheStoreURL == "" { + return 0 + } + return b.contextCacheDefaultTTL +} + // extProcUDSVolumeName is the name of the volume backing the extproc UDS socket, // shared between the extproc container and the envoy container. const extProcUDSVolumeName = mutationNamePrefix + "extproc-uds" diff --git a/internal/controller/extproc_builder_test.go b/internal/controller/extproc_builder_test.go index 8d7288370d..f818b72c29 100644 --- a/internal/controller/extproc_builder_test.go +++ b/internal/controller/extproc_builder_test.go @@ -8,6 +8,7 @@ package controller import ( "reflect" "testing" + "time" egv1a1 "github.com/envoyproxy/gateway/api/v1alpha1" "github.com/stretchr/testify/require" @@ -79,7 +80,7 @@ func TestExtProcContainerHash_Drift(t *testing.T) { input := extProcContainerInput{} baseHash := base.extProcContainerHash(input) require.NotEmpty(t, baseHash) - require.Equal(t, 20, reflect.TypeOf(extProcBuilder{}).NumField(), + require.Equal(t, 22, reflect.TypeOf(extProcBuilder{}).NumField(), "update drift coverage when extProcBuilder fields change") // Helper: mutate a copy of the builder, recompute, expect a different hash. @@ -91,6 +92,17 @@ func TestExtProcContainerHash_Drift(t *testing.T) { }) } + assertDrifts("contextCacheStoreURL", func() *extProcBuilder { + b := newTestBuilder() + b.contextCacheStoreURL = "redis:6379" + return b + }()) + assertDrifts("contextCacheDefaultTTL", func() *extProcBuilder { + b := newTestBuilder() + b.contextCacheStoreURL = "redis:6379" + b.contextCacheDefaultTTL = 10 * time.Minute + return b + }()) assertDrifts("extraEnvVars", func() *extProcBuilder { b := newTestBuilder() b.extraEnvVars = []corev1.EnvVar{{Name: "OTEL_SERVICE_NAME", Value: "ai-gateway"}} @@ -374,3 +386,15 @@ func TestImageTagOrDigest(t *testing.T) { }) } } + +// Context-cache args are passed only when a Redis URL is set, so deployments without +// caching keep identical args. +func TestBuildExtProcBaseArgs_ContextCache(t *testing.T) { + b := newTestBuilder() + require.NotContains(t, b.buildExtProcBaseArgs(false), "-contextCacheStoreURL") + + b.contextCacheStoreURL = "redis:6379" + b.contextCacheDefaultTTL = 10 * time.Minute + args := b.buildExtProcBaseArgs(false) + require.Subset(t, args, []string{"-contextCacheStoreURL", "redis:6379", "-contextCacheDefaultTTL", "10m0s"}) +} diff --git a/internal/controller/gateway.go b/internal/controller/gateway.go index f721a43aa1..2277419539 100644 --- a/internal/controller/gateway.go +++ b/internal/controller/gateway.go @@ -553,6 +553,7 @@ func (c *GatewayController) reconcileFilterConfigSecret( b.Schema = schemaToFilterAPI(backendObj.Spec.APISchema) b.HeaderValueFilters = headerValueFiltersToFilterAPI(backendObj.Spec.HeaderValueFilters) + } if bsp != nil { diff --git a/internal/extproc/processor_impl.go b/internal/extproc/processor_impl.go index 7883243ffb..9378cee34d 100644 --- a/internal/extproc/processor_impl.go +++ b/internal/extproc/processor_impl.go @@ -24,8 +24,10 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/protobuf/types/known/structpb" + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" "github.com/envoyproxy/ai-gateway/internal/backendauth" "github.com/envoyproxy/ai-gateway/internal/bodymutator" + "github.com/envoyproxy/ai-gateway/internal/contextcache" "github.com/envoyproxy/ai-gateway/internal/endpointspec" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/headermutator" @@ -33,6 +35,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/json" "github.com/envoyproxy/ai-gateway/internal/llmcostcel" "github.com/envoyproxy/ai-gateway/internal/metrics" + "github.com/envoyproxy/ai-gateway/internal/redaction" "github.com/envoyproxy/ai-gateway/internal/tracing/tracingapi" "github.com/envoyproxy/ai-gateway/internal/translator" ) @@ -129,6 +132,9 @@ type ( backendName string routeName string handler filterapi.BackendAuthHandler + // cacheResolver is the context-cache resolver for this backend. + // Non-nil only for GCP Vertex AI backends; set in SetBackend. + cacheResolver contextcache.CacheResolver // unsupportedBackendErr is set by SetBackend and answered as a 422 in ProcessRequestHeaders. unsupportedBackendErr error // cost is the cost of the request that is accumulated during the processing of the response. @@ -371,7 +377,47 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque // * The request is a streaming request, and the IncludeUsage option is set to false since we need to ensure that // the token usage is calculated correctly without being bypassed. forceBodyMutation := u.onRetry() || u.parent.forceBodyMutation + + // Resolve GCP context cache when the resolver is present and the translator supports + // cache injection. This must happen before RequestBody so the translator can use + // the resolved cache name and filtered message list. + var cacheApplied bool + var cacheName string + if u.cacheResolver != nil { + if cacheSetter, ok := u.translator.(translator.ContextCacheSetter); ok { + if req, ok := any(u.parent.originalRequestBody).(*openai.ChatCompletionRequest); ok { + result, resolveErr := u.cacheResolver.Resolve(ctx, req) + if resolveErr != nil { + u.logger.Error("context cache resolution failed", slog.String("error", resolveErr.Error())) + u.metrics.RecordRequestCompletion(ctx, false, u.requestHeaders) + return createUserFacingErrorResponse(502, "UpstreamError", "upstream cache service error"), nil + } + if result != nil { + cacheResult := &translator.ContextCacheResult{ + CacheName: result.CacheName, + FilteredMessages: result.Messages, + Created: result.Created, + } + // Only a create is billed as a cache write; a reused cache is not. + if result.Created { + cacheResult.WriteTokenCount = uint32(result.TokenCount) //nolint:gosec + } + cacheSetter.SetContextCacheResult(cacheResult) + cacheApplied, cacheName = true, result.CacheName + } + } + } + } + newHeaders, newBody, err := u.translator.RequestBody(u.parent.originalRequestBodyRaw, u.parent.originalRequestBody, forceBodyMutation) + if err == nil && cacheApplied && u.parent.debugLogEnabled { + body := string(newBody) + if u.parent.enableRedaction { + body = redaction.RedactString(body) + } + u.logger.Debug("request body rewritten for context cache", + slog.String("cache", cacheName), slog.String("body", body)) + } if err != nil { if userFacingErr := internalapi.GetUserFacingError(err); userFacingErr != nil { // return to user as 422 - e.g., "invalid request body: tool_choice type not supported" @@ -752,6 +798,7 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) SetBackend(c u.backendName = backend.Backend.Name u.routeName = routeName u.handler = backend.Handler + u.cacheResolver = backend.CacheResolver u.headerMutator = headermutator.NewHeaderMutator(backend.Backend.HeaderMutation, rp.requestHeaders) u.bodyMutator = bodymutator.NewBodyMutator(backend.Backend.BodyMutation, rp.originalRequestBodyRaw) // Header-derived labels/CEL must be able to see the overridden request model. diff --git a/internal/extproc/processor_impl_test.go b/internal/extproc/processor_impl_test.go index 441edae767..38482a780b 100644 --- a/internal/extproc/processor_impl_test.go +++ b/internal/extproc/processor_impl_test.go @@ -24,12 +24,14 @@ import ( "github.com/google/cel-go/cel" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/propagation" + "golang.org/x/oauth2" "google.golang.org/protobuf/types/known/structpb" anthropicschema "github.com/envoyproxy/ai-gateway/internal/apischema/anthropic" "github.com/envoyproxy/ai-gateway/internal/apischema/openai" "github.com/envoyproxy/ai-gateway/internal/backendauth" "github.com/envoyproxy/ai-gateway/internal/bodymutator" + "github.com/envoyproxy/ai-gateway/internal/contextcache" "github.com/envoyproxy/ai-gateway/internal/endpointspec" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/headermutator" @@ -2599,3 +2601,240 @@ func mustCompileCEL(t *testing.T, expr string) cel.Program { require.NoError(t, err) return prog } + +// --------------------------------------------------------------------------- +// GCP context-cache resolver wiring tests +// --------------------------------------------------------------------------- + +// mockCacheResolver is a controllable contextcache.CacheResolver for testing. +type mockCacheResolver struct { + retResult *contextcache.ResolveResult + retErr error + called bool +} + +func (m *mockCacheResolver) Resolve(_ context.Context, _ *openai.ChatCompletionRequest) (*contextcache.ResolveResult, error) { + m.called = true + return m.retResult, m.retErr +} + +// mockCacheTranslator extends mockTranslator with ContextCacheSetter support. +type mockCacheTranslator struct { + mockTranslator + setCacheResult *translator.ContextCacheResult +} + +func (m *mockCacheTranslator) SetContextCacheResult(result *translator.ContextCacheResult) { + m.setCacheResult = result +} + +// mockGCPAuthHandler implements filterapi.GCPAuthHandler for testing. +type mockGCPAuthHandler struct{} + +func (m *mockGCPAuthHandler) Do(_ context.Context, _ map[string]string, _ []byte) ([]internalapi.Header, error) { + return nil, nil +} + +func (m *mockGCPAuthHandler) GCPTokenSource() oauth2.TokenSource { + return oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "test-token"}) +} + +func (m *mockGCPAuthHandler) GCPRegion() string { return "us-central1" } +func (m *mockGCPAuthHandler) GCPProject() string { return "test-project" } + +// Test_upstreamProcessor_CacheResolver_notCalled verifies that the resolver is not invoked +// when the translator does not implement ContextCacheSetter. +func Test_upstreamProcessor_CacheResolver_notCalled(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockCacheResolver{} + // mockTranslator does not implement ContextCacheSetter. + tr := &mockTranslator{t: t, expRequestBody: &body} + mm := &mockMetrics{} + + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, + logger: slog.Default(), + originalRequestBodyRaw: someBody, + originalRequestBody: &body, + originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: mm, + translator: tr, + handler: &mockGCPAuthHandler{}, + cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.False(t, cr.called, "resolver must not be called when translator lacks ContextCacheSetter") +} + +// Test_upstreamProcessor_CacheResolver_noMarkers verifies that when the resolver returns nil +// (no cache_control markers), the translator is called without a cache result. +func Test_upstreamProcessor_CacheResolver_noMarkers(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockCacheResolver{retResult: nil} // no markers + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + mm := &mockMetrics{} + + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, + logger: slog.Default(), + originalRequestBodyRaw: someBody, + originalRequestBody: &body, + originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: mm, + translator: tr, + handler: &mockGCPAuthHandler{}, + cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.True(t, cr.called) + require.Nil(t, tr.setCacheResult, "SetContextCacheResult must not be called when resolver returns nil") +} + +// Test_upstreamProcessor_CacheResolver_hit verifies that a cache hit is injected into the translator. +func Test_upstreamProcessor_CacheResolver_hit(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + filteredMessages := []openai.ChatCompletionMessageParamUnion{} + cr := &mockCacheResolver{retResult: &contextcache.ResolveResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/abc123", + Messages: filteredMessages, + Created: false, + }} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + mm := &mockMetrics{} + + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, + logger: slog.Default(), + originalRequestBodyRaw: someBody, + originalRequestBody: &body, + originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: mm, + translator: tr, + handler: &mockGCPAuthHandler{}, + cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.True(t, cr.called) + require.NotNil(t, tr.setCacheResult) + require.Equal(t, "projects/p/locations/us-central1/cachedContents/abc123", tr.setCacheResult.CacheName) + require.Equal(t, filteredMessages, tr.setCacheResult.FilteredMessages) + require.False(t, tr.setCacheResult.Created) + require.Zero(t, tr.setCacheResult.WriteTokenCount) +} + +// Test_upstreamProcessor_CacheResolver_created verifies that cache-write tokens are propagated. +func Test_upstreamProcessor_CacheResolver_created(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockCacheResolver{retResult: &contextcache.ResolveResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/new456", + Messages: []openai.ChatCompletionMessageParamUnion{}, + Created: true, + TokenCount: 512, + }} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + mm := &mockMetrics{} + + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, + logger: slog.Default(), + originalRequestBodyRaw: someBody, + originalRequestBody: &body, + originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: mm, + translator: tr, + handler: &mockGCPAuthHandler{}, + cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.True(t, cr.called) + require.NotNil(t, tr.setCacheResult) + require.True(t, tr.setCacheResult.Created) + require.Equal(t, uint32(512), tr.setCacheResult.WriteTokenCount) +} + +// A hit must not be billed as a write even if the resolver reports a token count. +func Test_upstreamProcessor_CacheResolver_hitIgnoresTokenCount(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockCacheResolver{retResult: &contextcache.ResolveResult{ + CacheName: "c/hit", Messages: []openai.ChatCompletionMessageParamUnion{}, Created: false, TokenCount: 99, + }} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, logger: slog.Default(), + originalRequestBodyRaw: someBody, originalRequestBody: &body, originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: &mockMetrics{}, translator: tr, handler: &mockGCPAuthHandler{}, cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.NotNil(t, tr.setCacheResult) + require.Zero(t, tr.setCacheResult.WriteTokenCount) +} + +// Test_upstreamProcessor_CacheResolver_error verifies that a resolver error produces a 502. +func Test_upstreamProcessor_CacheResolver_error(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockCacheResolver{retErr: errors.New("GCP API quota exceeded")} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + mm := &mockMetrics{} + + p := &chatCompletionProcessorUpstreamFilter{ + parent: &chatCompletionProcessorRouterFilter{ + config: &filterapi.RuntimeConfig{}, + logger: slog.Default(), + originalRequestBodyRaw: someBody, + originalRequestBody: &body, + originalModel: "gemini-1.5-pro", + }, + requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, + metrics: mm, + translator: tr, + handler: &mockGCPAuthHandler{}, + cacheResolver: cr, + logger: slog.Default(), + } + resp, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.NotNil(t, resp) + immediateResp, ok := resp.Response.(*extprocv3.ProcessingResponse_ImmediateResponse) + require.True(t, ok) + require.Equal(t, typev3.StatusCode(502), immediateResp.ImmediateResponse.Status.Code) + require.Contains(t, string(immediateResp.ImmediateResponse.Body), "upstream cache service error") + mm.RequireRequestFailure(t) +} diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 1ef537cf2b..31c4988267 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -14,6 +14,7 @@ import ( "slices" "strings" "sync" + "time" "unicode/utf8" corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" @@ -26,6 +27,8 @@ import ( "google.golang.org/protobuf/types/known/structpb" "github.com/envoyproxy/ai-gateway/internal/backendauth" + "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/contextcache/gcpcache" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/redaction" @@ -61,6 +64,17 @@ type Server struct { routerProcessorsPerReqID map[string]Processor routerProcessorsPerReqIDMutex sync.RWMutex uuidFn func() string + // cacheStore is the shared context-cache store, set by EnableContextCache. Nil + // disables context caching for every backend. + cacheStore contextcache.Store + // cacheDefaultTTL is the cache TTL used when a cache_control marker gives none. + cacheDefaultTTL time.Duration + // cacheResolvers holds one context-cache resolver per GCP Vertex AI backend, keyed by + // backend name. Resolvers are carried across config reloads so that their background + // syncers keep running and are not restarted on every filterapi update. + cacheResolvers map[string]*gcpcache.Resolver + // cacheResolversMu guards cacheResolvers against LoadConfig racing Close. + cacheResolversMu sync.Mutex } // NewServer creates a new external processor server. @@ -72,21 +86,84 @@ func NewServer(logger *slog.Logger, enableRedaction bool) (*Server, error) { enableRedaction: enableRedaction, processorFactories: make(map[string]ProcessorFactory), routerProcessorsPerReqID: make(map[string]Processor), + cacheResolvers: make(map[string]*gcpcache.Resolver), uuidFn: uuid.NewString, } return srv, nil } +// EnableContextCache turns on context caching for every GCP Vertex AI backend, using store +// to share resolved caches across replicas. It must be called before the first LoadConfig. +func (s *Server) EnableContextCache(store contextcache.Store, defaultTTL time.Duration) { + s.cacheStore = store + s.cacheDefaultTTL = defaultTTL +} + // LoadConfig updates the configuration of the external processor. func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error { newConfig, err := filterapi.NewRuntimeConfig(ctx, config, backendauth.NewHandler) if err != nil { return fmt.Errorf("cannot create runtime filter config: %w", err) } + + // LoadConfig is reachable on a Server built directly rather than through NewServer + // (notably in tests), so the logger is resolved here rather than assumed non-nil. + logger := s.logger + if logger == nil { + logger = slog.New(slog.DiscardHandler) + } + + // Attach a context-cache resolver to each GCP Vertex AI (Gemini) backend when context + // caching is enabled. Only the Gemini translator uses resolved caches, so other GCP + // schemas (such as GCPAnthropic) get none. Resolvers are keyed by backend name and + // reused across reloads; a reused resolver is given the reload's auth handler, since the + // controller rotates the token. Resolvers whose backend disappeared or changed schema + // are closed, which stops their background syncer. + s.cacheResolversMu.Lock() + defer s.cacheResolversMu.Unlock() + live := make(map[string]*gcpcache.Resolver, len(s.cacheResolvers)) + for _, rb := range newConfig.Backends { + if s.cacheStore == nil || rb.Backend.Schema.Name != filterapi.APISchemaGCPVertexAI { + continue + } + auth, ok := rb.Handler.(filterapi.GCPAuthHandler) + if !ok { + continue + } + name := rb.Backend.Name + resolver, ok := s.cacheResolvers[name] + if ok { + resolver.SetAuth(auth) + } else { + resolver = gcpcache.New(nil, s.cacheStore, auth, s.cacheDefaultTTL, logger) + } + // Start is a no-op for a resolver already running. It runs under the resolver's own + // context, not ctx, which is cancelled when this reload finishes. + resolver.Start() + live[name] = resolver + rb.CacheResolver = resolver + } + for name, resolver := range s.cacheResolvers { + if _, ok := live[name]; !ok { + _ = resolver.Close() + } + } + s.cacheResolvers = live + s.config = newConfig // This is racey, but we don't care. return nil } +// Close stops the background syncer of every context-cache resolver. It is safe to call +// more than once. +func (s *Server) Close() { + s.cacheResolversMu.Lock() + defer s.cacheResolversMu.Unlock() + for _, resolver := range s.cacheResolvers { + _ = resolver.Close() + } +} + // Register a new processor for the given request path. func (s *Server) Register(path string, newProcessor ProcessorFactory) { s.logger.Info("Registering processor", slog.String("path", path)) diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index 8ae104251a..f0d17c6bd0 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -19,11 +19,13 @@ import ( extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3" "github.com/stretchr/testify/require" + "go.uber.org/goleak" "google.golang.org/grpc/codes" "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/structpb" + "github.com/envoyproxy/ai-gateway/internal/contextcache" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/internalapi" internaltesting "github.com/envoyproxy/ai-gateway/internal/testing" @@ -48,11 +50,105 @@ func requireNewServerWithMockProcessor(t *testing.T) (*Server, *mockProcessor) { func TestServer_LoadConfig(t *testing.T) { config := &filterapi.Config{} s := &Server{} + t.Cleanup(s.Close) err := s.LoadConfig(t.Context(), config) require.NoError(t, err) require.NotNil(t, s.config) } +// cachingServer returns a Server with context caching enabled against an in-memory store. +func cachingServer(t *testing.T) *Server { + t.Helper() + s := &Server{} + s.EnableContextCache(contextcache.NoopStore{}, time.Minute) + t.Cleanup(s.Close) + return s +} + +// gcpBackend builds a GCP Vertex AI backend with GCP credentials. +func gcpBackend(name string) filterapi.Backend { + return filterapi.Backend{ + Name: name, + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}}, + } +} + +// A resolver is attached only to GCP Vertex AI backends with GCP credentials. +func TestServer_LoadConfig_CacheResolver(t *testing.T) { + s := cachingServer(t) + anthropic := gcpBackend("gcp-anthropic") + anthropic.Schema = filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPAnthropic} + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackend("gemini"), + anthropic, + {Name: "openai", Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}}, + }})) + + require.NotNil(t, s.config.Backends["gemini"].CacheResolver) + require.Nil(t, s.config.Backends["gcp-anthropic"].CacheResolver, "only Gemini uses resolved caches") + require.Nil(t, s.config.Backends["openai"].CacheResolver) +} + +// Without EnableContextCache no backend gets a resolver. +func TestServer_LoadConfig_CacheDisabled(t *testing.T) { + s := &Server{} + t.Cleanup(s.Close) + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{gcpBackend("gemini")}})) + require.Nil(t, s.config.Backends["gemini"].CacheResolver) + require.Empty(t, s.cacheResolvers) +} + +// An unchanged backend keeps its resolver across reloads, and gets the reload's new +// handler, since the controller rotates the access token. +func TestServer_LoadConfig_CacheResolver_ReusedWithNewHandler(t *testing.T) { + s := cachingServer(t) + config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackend("gemini")}} + + require.NoError(t, s.LoadConfig(t.Context(), config)) + first := s.config.Backends["gemini"].CacheResolver + firstHandler := s.config.Backends["gemini"].Handler + + require.NoError(t, s.LoadConfig(t.Context(), config)) + require.Same(t, first, s.config.Backends["gemini"].CacheResolver) + second := s.config.Backends["gemini"].Handler + require.NotSame(t, firstHandler, second, "the handler is rebuilt on every reload") + require.Equal(t, second, s.cacheResolvers["gemini"].GetAuth()) +} + +// A resolver evicted by a reload must have its background syncer stopped. +func TestServer_LoadConfig_CacheResolver_EvictedIsClosed(t *testing.T) { + before := goleak.IgnoreCurrent() + s := &Server{} + s.EnableContextCache(contextcache.NoopStore{}, time.Minute) + + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackend("a"), gcpBackend("b"), + }})) + kept := s.cacheResolvers["a"] + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{gcpBackend("a")}})) + require.Len(t, s.cacheResolvers, 1) + // If the evicted resolver had not been closed, its syncer would still be running. + require.NoError(t, kept.Close()) + goleak.VerifyNone(t, before) +} + +// The reload context is cancelled when the reload finishes; syncers must outlive it. +func TestServer_LoadConfig_SyncerOutlivesReloadContext(t *testing.T) { + before := goleak.IgnoreCurrent() + s := &Server{} + s.EnableContextCache(contextcache.NoopStore{}, time.Minute) + + ctx, cancel := context.WithCancel(t.Context()) + require.NoError(t, s.LoadConfig(ctx, &filterapi.Config{Backends: []filterapi.Backend{gcpBackend("gemini")}})) + cancel() + time.Sleep(50 * time.Millisecond) + require.Error(t, goleak.Find(before), "the syncer must still be running after the reload context ends") + + s.Close() + goleak.VerifyNone(t, before) +} + func TestServer_Check(t *testing.T) { s, _ := requireNewServerWithMockProcessor(t) diff --git a/internal/filterapi/runtime.go b/internal/filterapi/runtime.go index 3885fb71f9..695a8ee54d 100644 --- a/internal/filterapi/runtime.go +++ b/internal/filterapi/runtime.go @@ -10,7 +10,9 @@ import ( "fmt" "github.com/google/cel-go/cel" + "golang.org/x/oauth2" + "github.com/envoyproxy/ai-gateway/internal/contextcache" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/llmcostcel" ) @@ -22,6 +24,20 @@ type BackendAuthHandler interface { Do(ctx context.Context, requestHeaders map[string]string, mutatedBody []byte) ([]internalapi.Header, error) } +// GCPAuthHandler is an optional interface implemented by the GCP backend auth handler. +// It exposes credentials so that in-process components (e.g. the context-cache resolver) can +// authenticate to GCP APIs without duplicating credential management. +type GCPAuthHandler interface { + BackendAuthHandler + // GCPTokenSource returns the OAuth2 token source used to authenticate requests to GCP. + // Callers must call Token() on the returned source to obtain a valid access token. + GCPTokenSource() oauth2.TokenSource + // GCPRegion returns the GCP region configured for this backend. + GCPRegion() string + // GCPProject returns the GCP project name configured for this backend. + GCPProject() string +} + // NewBackendAuthHandlerFunc is a function type that creates a new BackendAuthHandler for a given BackendAuth configuration. type NewBackendAuthHandlerFunc func(ctx context.Context, auth *BackendAuth) (BackendAuthHandler, error) @@ -53,6 +69,9 @@ type RuntimeBackend struct { Backend *Backend // Handler is the backend auth handler. Handler BackendAuthHandler + // CacheResolver is the in-process context-cache resolver for this backend. Currently + // only used for GCP Vertex AI backends with context caching configured. + CacheResolver contextcache.CacheResolver } // RuntimeGlobalRequestCost is the configuration for gateway-level default request costs. diff --git a/internal/json/json.go b/internal/json/json.go index 2cb1caf187..e3da765f6a 100644 --- a/internal/json/json.go +++ b/internal/json/json.go @@ -23,7 +23,15 @@ var ( // NewEncoder is equivalent to encoding/json.NewEncoder. NewEncoder = config.NewEncoder // NewDecoder is equivalent to encoding/json.NewDecoder. - NewDecoder = config.NewDecoder + NewDecoder = config.NewDecoder + sortedConfig = sonicjson.Config{ + CaseSensitive: true, + SortMapKeys: true, + }.Froze() + + // MarshalSorted is Marshal with map keys sorted, so equal values always produce the + // same bytes. Use it when the output is hashed or compared, not on hot paths. + MarshalSorted = sortedConfig.Marshal // MarshalForDeterministicTesting marshals a value to JSON in a deterministic way for testing. // The normal sonic configuration does not guarantee deterministic output in terms of field order. // It panics if called outside of tests. diff --git a/internal/translator/gemini_helper.go b/internal/translator/gemini_helper.go index 9d9178f2b5..ad964a9b2a 100644 --- a/internal/translator/gemini_helper.go +++ b/internal/translator/gemini_helper.go @@ -53,6 +53,19 @@ const ( // Request Conversion Helper for OpenAI to GCP Gemini Translator // -------------------------------------------------------------. +// OpenAIMessagesToGeminiContents converts OpenAI messages to Gemini Contents and SystemInstruction. +// It is exported so that the gcpcache resolver can build the cached prefix in Gemini format +// without duplicating conversion logic. +func OpenAIMessagesToGeminiContents(messages []openai.ChatCompletionMessageParamUnion, requestModel internalapi.RequestModel) ([]genai.Content, *genai.Content, error) { + return openAIMessagesToGeminiContents(messages, requestModel) +} + +// OpenAIToolsToGeminiTools converts OpenAI tools to Gemini tools. +// It is exported so that the gcpcache resolver can include tools in the cached prefix. +func OpenAIToolsToGeminiTools(openaiTools []openai.Tool, parametersJSONSchemaAvailable bool) ([]genai.Tool, error) { + return openAIToolsToGeminiTools(openaiTools, parametersJSONSchemaAvailable) +} + // openAIMessagesToGeminiContents converts OpenAI messages to Gemini Contents and SystemInstruction. func openAIMessagesToGeminiContents(messages []openai.ChatCompletionMessageParamUnion, requestModel internalapi.RequestModel) ([]genai.Content, *genai.Content, error) { var gcpContents []genai.Content diff --git a/internal/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index 529f265ab8..677b746335 100644 --- a/internal/translator/openai_gcpvertexai.go +++ b/internal/translator/openai_gcpvertexai.go @@ -92,6 +92,14 @@ type openAIToGCPVertexAITranslatorV1ChatCompletion struct { debugLogEnabled bool enableRedaction bool logger *slog.Logger + // cacheResult is set by SetContextCacheResult before RequestBody is called. + // When non-nil, the cache name is injected into the Gemini request and the + // filtered (non-cached) messages replace the original message list. + cacheResult *ContextCacheResult + // cacheWriteTokens is the number of tokens stored in a newly created cache entry. + // Set during RequestBody (from cacheResult) and consumed in ResponseBody to + // record cache-write cost via tokenUsage.SetCacheCreationInputTokens. + cacheWriteTokens uint32 } // RequestBody implements [OpenAIChatCompletionTranslator.RequestBody] for GCP Gemini. @@ -120,6 +128,13 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) RequestBody(_ []byte, op if err != nil { return nil, nil, err } + + // Capture cache-write token count for cost attribution in ResponseBody. Only a call + // that created the cache pays for the write. + if o.cacheResult != nil && o.cacheResult.Created { + o.cacheWriteTokens = o.cacheResult.WriteTokenCount + } + newBody, err = json.Marshal(gcpReq) if err != nil { return nil, nil, fmt.Errorf("error marshaling Gemini request: %w", err) @@ -196,6 +211,10 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseBody(_ map[strin if openAIResp.Usage.CompletionTokensDetails != nil { tokenUsage.SetReasoningTokens(uint32(openAIResp.Usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec } + // Attribute cache-write tokens when a new cache entry was created for this request. + if o.cacheWriteTokens > 0 { + tokenUsage.SetCacheCreationInputTokens(o.cacheWriteTokens) + } if span != nil { span.RecordResponse(openAIResp) @@ -260,6 +279,10 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) handleStreamingResponse( tokenUsage.SetTotalTokens(uint32(usage.TotalTokens)) //nolint:gosec tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens)) //nolint:gosec tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec + // Attribute cache-write tokens when a new cache entry was created for this request. + if o.cacheWriteTokens > 0 { + tokenUsage.SetCacheCreationInputTokens(o.cacheWriteTokens) + } } } @@ -517,8 +540,15 @@ func getGenerationConfigThinkingConfig(tu *openai.ThinkingUnion) *genai.Thinking // openAIMessageToGeminiMessage converts an OpenAI ChatCompletionRequest to a GCP Gemini GenerateContentRequest. func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMessage(openAIReq *openai.ChatCompletionRequest, requestModel internalapi.RequestModel) (*gcp.GenerateContentRequest, error) { + // If a cache result was set (by SetContextCacheResult), use the filtered (non-cached) + // messages so that only the new turns are sent in the Gemini request body. + messages := openAIReq.Messages + if o.cacheResult != nil && o.cacheResult.FilteredMessages != nil { + messages = o.cacheResult.FilteredMessages + } + // Convert OpenAI messages to Gemini Contents and SystemInstruction. - contents, systemInstruction, err := openAIMessagesToGeminiContents(openAIReq.Messages, requestModel) + contents, systemInstruction, err := openAIMessagesToGeminiContents(messages, requestModel) if err != nil { return nil, err } @@ -559,6 +589,16 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes // Vendor fields take precedence over translated fields when conflicts occur. o.applyVendorSpecificFields(openAIReq, &gcr, requestModel) + // Inject resolved cache name from the context-cache resolver (set via SetContextCacheResult). + if o.cacheResult != nil { + gcr.CachedContent = o.cacheResult.CacheName + // Vertex rejects cached_content combined with tools, tool_config or + // system_instruction; the tools and system prompt already live in the cache. + gcr.Tools = nil + gcr.ToolConfig = nil + gcr.SystemInstruction = nil + } + return &gcr, nil } @@ -586,6 +626,14 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) applyVendorSpecificField } } +// SetContextCacheResult implements [ContextCacheSetter]. The upstream processor calls this +// before RequestBody when the context-cache resolver has resolved or created a cache entry. +// The stored result is consumed in openAIMessageToGeminiMessage (to inject cachedContent +// and replace the message list) and cleared after use so retries start fresh. +func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) SetContextCacheResult(result *ContextCacheResult) { + o.cacheResult = result +} + func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) geminiResponseToOpenAIMessage(gcr *genai.GenerateContentResponse, responseModel string) (*openai.ChatCompletionResponse, error) { // Convert candidates to OpenAI choices. choices, err := geminiCandidatesToOpenAIChoices(gcr.Candidates, o.responseMode) diff --git a/internal/translator/openai_gcpvertexai_test.go b/internal/translator/openai_gcpvertexai_test.go index 439c82d219..75a7091584 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -2742,3 +2742,145 @@ func TestGCPVertexAIRedactBody(t *testing.T) { require.NotContains(t, *resp.Choices[0].Message.Content, "[REDACTED") }) } + +// TestOpenAIToGCPVertexAITranslator_ContextCacheSetter verifies that SetContextCacheResult +// correctly injects the cache name into the Gemini request, replaces the message list +// with the filtered remainder, and records cache-write tokens in ResponseBody. +func TestOpenAIToGCPVertexAITranslator_ContextCacheSetter(t *testing.T) { + t.Run("cache name injected into Gemini request", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + + // Seed a cache result before RequestBody. + tr.SetContextCacheResult(&ContextCacheResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/abc", + FilteredMessages: []openai.ChatCompletionMessageParamUnion{}, + Created: true, + WriteTokenCount: 256, + }) + + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfUser: &openai.ChatCompletionUserMessageParam{ + Content: openai.StringOrUserRoleContentUnion{Value: "hello"}, + Role: openai.ChatMessageRoleUser, + }}, + }, + } + + _, bodyBytes, err := tr.RequestBody(nil, req, false) + require.NoError(t, err) + + // The Gemini request should carry the cachedContent field. + require.Contains(t, string(bodyBytes), `"cachedContent":"projects/p/locations/us-central1/cachedContents/abc"`) + + // cacheWriteTokens should be set for response attribution. + require.Equal(t, uint32(256), tr.cacheWriteTokens) + }) + + t.Run("cached request drops tools, tool config and system instruction", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + tr.SetContextCacheResult(&ContextCacheResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/abc", + FilteredMessages: []openai.ChatCompletionMessageParamUnion{ + {OfSystem: &openai.ChatCompletionSystemMessageParam{ + Content: openai.ContentUnion{Value: "late system"}, + Role: openai.ChatMessageRoleSystem, + }}, + {OfUser: &openai.ChatCompletionUserMessageParam{ + Content: openai.StringOrUserRoleContentUnion{Value: "hello"}, + Role: openai.ChatMessageRoleUser, + }}, + }, + }) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Tools: []openai.Tool{{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{Name: "get_weather"}, + }}, + ToolChoice: &openai.ChatCompletionToolChoiceUnion{Value: "auto"}, + } + _, bodyBytes, err := tr.RequestBody(nil, req, false) + require.NoError(t, err) + body := string(bodyBytes) + require.Contains(t, body, `"cachedContent"`) + require.NotContains(t, body, "get_weather") + require.NotContains(t, body, `"toolConfig":{`) + require.NotContains(t, body, "late system") + require.Contains(t, body, "hello") + }) + + t.Run("cache hit does not set write tokens", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + + tr.SetContextCacheResult(&ContextCacheResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/existing", + Created: false, + WriteTokenCount: 0, + }) + + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{}, + } + _, _, err := tr.RequestBody(nil, req, false) + require.NoError(t, err) + require.Equal(t, uint32(0), tr.cacheWriteTokens) + }) + + t.Run("a hit reports no write tokens even if a count is set", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + tr.SetContextCacheResult(&ContextCacheResult{CacheName: "c/hit", Created: false, WriteTokenCount: 99}) + _, _, err := tr.RequestBody(nil, &openai.ChatCompletionRequest{Model: "gemini-1.5-pro"}, false) + require.NoError(t, err) + require.Zero(t, tr.cacheWriteTokens) + }) + + t.Run("cache write tokens appear in non-streaming ResponseBody", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + tr.cacheWriteTokens = 512 + + gcpResp := `{ + "candidates": [{ + "content": {"parts": [{"text": "hello"}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 100, + "candidatesTokenCount": 20, + "totalTokenCount": 120 + } + }` + + _, _, tokenUsage, _, err := tr.ResponseBody(nil, strings.NewReader(gcpResp), true, nil) + require.NoError(t, err) + + cacheCreation, set := tokenUsage.CacheCreationInputTokens() + require.True(t, set, "CacheCreationInputTokens must be set when cacheWriteTokens > 0") + require.Equal(t, uint32(512), cacheCreation) + }) + + t.Run("no cache write tokens when not created", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + // cacheWriteTokens defaults to zero. + + gcpResp := `{ + "candidates": [{ + "content": {"parts": [{"text": "hello"}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 50, + "candidatesTokenCount": 10, + "totalTokenCount": 60 + } + }` + + _, _, tokenUsage, _, err := tr.ResponseBody(nil, strings.NewReader(gcpResp), true, nil) + require.NoError(t, err) + + _, set := tokenUsage.CacheCreationInputTokens() + require.False(t, set, "CacheCreationInputTokens must not be set when no cache was created") + }) +} diff --git a/internal/translator/translator.go b/internal/translator/translator.go index b79a18de84..e84abc6eaa 100644 --- a/internal/translator/translator.go +++ b/internal/translator/translator.go @@ -138,6 +138,31 @@ type ResponseRedactor interface { RedactBody(resp *openai.ChatCompletionResponse) *openai.ChatCompletionResponse } +// ContextCacheSetter is an optional interface implemented by translators whose provider +// supports context caching (currently GCP Vertex AI Gemini). The upstream processor calls +// SetContextCacheResult after the context-cache resolver runs so that the translator can +// (a) reference the provider cache in the upstream request and (b) record cache-write +// token counts for cost attribution in the response phase. +type ContextCacheSetter interface { + // SetContextCacheResult forwards the resolved (or newly created) cache entry to the translator. + // Must be called before RequestBody. + SetContextCacheResult(result *ContextCacheResult) +} + +// ContextCacheResult carries the information returned by the cache resolver that the translator needs. +type ContextCacheResult struct { + // CacheName is the provider's name for the resolved or created cache entry. + CacheName string + // FilteredMessages is the non-cached remainder of the conversation (messages after the breakpoint). + // The translator replaces the original messages with this slice when building the provider request. + FilteredMessages []openai.ChatCompletionMessageParamUnion + // Created is true when this call created a new cache entry (cache-write cost applies). + Created bool + // WriteTokenCount is the number of tokens stored in the cache, from the provider's create response. + // Only meaningful when Created is true. + WriteTokenCount uint32 +} + // AnthropicResponseRedactor is an optional interface that Anthropic translators // can implement to support response body redaction for debug logging. type AnthropicResponseRedactor interface { diff --git a/manifests/charts/ai-gateway-helm/templates/deployment.yaml b/manifests/charts/ai-gateway-helm/templates/deployment.yaml index c070704e3b..543bdb597e 100644 --- a/manifests/charts/ai-gateway-helm/templates/deployment.yaml +++ b/manifests/charts/ai-gateway-helm/templates/deployment.yaml @@ -62,6 +62,12 @@ spec: {{- if .Values.extProc.enableRedaction }} - --extProcEnableRedaction=true {{- end }} + {{- with .Values.extProc.contextCache }} + {{- if .storeURL }} + - --extProcContextCacheStoreURL={{ .storeURL }} + - --extProcContextCacheDefaultTTL={{ .defaultTTL | default "300s" }} + {{- end }} + {{- end }} {{- with (include "ai-gateway-helm.extProc.imagePullSecretsString" .) }} - --extProcImagePullSecrets={{ . }} {{- end }} diff --git a/manifests/charts/ai-gateway-helm/values.yaml b/manifests/charts/ai-gateway-helm/values.yaml index 16b5d32ada..28975dfc86 100644 --- a/manifests/charts/ai-gateway-helm/values.yaml +++ b/manifests/charts/ai-gateway-helm/values.yaml @@ -52,6 +52,15 @@ extProc: logFormat: text # Enable redaction of sensitive information in debug logs (only takes effect when logLevel is "debug"). enableRedaction: false + # Context caching for GCP Vertex AI (Gemini) backends. Setting storeURL enables it for every + # GCP Vertex AI backend; leaving it empty disables it. + contextCache: + # Store shared by all extproc replicas, as "host:port" or a "redis://" URL of a + # Redis-compatible server. A URL may contain credentials, which end up in the extproc + # container args. + storeURL: "" + # Cache TTL used when a cache_control marker does not specify one. Whole seconds, at least 60s. + defaultTTL: 300s ## @param extProc.extraEnvVars Array with extra environment variables to add to extProc containers ## e.g: