From c056b8763a76892bdeb284b21719da7ad5f7583a Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Fri, 31 Jul 2026 12:52:05 +0530 Subject: [PATCH 01/17] translator/gcpvertexai: add explicit cachedContent passthrough Add support for the GCP Vertex AI context caching explicit passthrough path: - Add CachedContent field to gcp.GenerateContentRequest wire schema so the Gemini API receives the cache reference. - Add CachedContent field to GCPVertexAIVendorFields (inline on ChatCompletionRequest) so callers can supply a pre-created cache resource name via the OpenAI-compatible API. - Copy the vendor field into the Gemini request in applyVendorSpecificFields. - Validate that cache_control markers and explicit cachedContent are not used together; return ErrMalformedRequest (HTTP 400) when both are present. - Add gcpRequestHasCacheControlMarkers helper to detect Anthropic-style cache_control on system/user/tool message content parts. - Add unit tests covering passthrough, validation error, and marker detection across all message types. Signed-off-by: Sivanantham Chinnaiyan --- internal/apischema/gcp/gcp.go | 5 + internal/apischema/openai/openai.go | 9 + internal/translator/openai_gcpvertexai.go | 51 +++++ .../translator/openai_gcpvertexai_test.go | 179 ++++++++++++++++++ 4 files changed, 244 insertions(+) diff --git a/internal/apischema/gcp/gcp.go b/internal/apischema/gcp/gcp.go index d2bff89f62..d38ed0d9b0 100644 --- a/internal/apischema/gcp/gcp.go +++ b/internal/apischema/gcp/gcp.go @@ -39,6 +39,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://cloud.google.com/vertex-ai/docs/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 diff --git a/internal/apischema/openai/openai.go b/internal/apischema/openai/openai.go index 785629ae30..94bf98908d 100644 --- a/internal/apischema/openai/openai.go +++ b/internal/apischema/openai/openai.go @@ -2042,6 +2042,15 @@ type GCPVertexAIVendorFields struct { // // https://cloud.google.com/vertex-ai/docs/reference/rest/v1/SafetySetting SafetySettings []*genai.SafetySetting `json:"safetySettings,omitzero"` + + // CachedContent specifies a pre-existing Vertex AI cached content resource to use as context. + // Format: "projects/{project}/locations/{location}/cachedContents/{cache_id}" + // + // When provided, cache_control markers on messages must be absent. The cache reference is + // sent directly to Gemini without calling any cache resolution logic. + // + // https://cloud.google.com/vertex-ai/docs/context-cache/context-cache-overview + CachedContent string `json:"cachedContent,omitzero"` } // GCPVertexAIGenerationConfig represents Gemini generation configuration options. diff --git a/internal/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index 529f265ab8..b1039031e4 100644 --- a/internal/translator/openai_gcpvertexai.go +++ b/internal/translator/openai_gcpvertexai.go @@ -105,6 +105,12 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) RequestBody(_ []byte, op o.requestModel = o.modelNameOverride } + // Validate that explicit cachedContent and cache_control markers are not used together. + if openAIReq.GCPVertexAIVendorFields != nil && openAIReq.CachedContent != "" && + gcpRequestHasCacheControlMarkers(openAIReq) { + return nil, nil, fmt.Errorf("%w: cannot specify both cache_control on messages and explicit cachedContent field", internalapi.ErrMalformedRequest) + } + // Set streaming flag. o.stream = openAIReq.Stream @@ -562,6 +568,48 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes return &gcr, nil } +// gcpRequestHasCacheControlMarkers returns true if any message content part in the request +// contains a cache_control marker (Anthropic-style ephemeral cache breakpoint). +func gcpRequestHasCacheControlMarkers(openAIReq *openai.ChatCompletionRequest) bool { + for i := range openAIReq.Messages { + msg := &openAIReq.Messages[i] + // Check tool messages (AnthropicContentFields inline on message itself). + if msg.OfTool != nil && msg.OfTool.AnthropicContentFields != nil && + isCacheEnabled(msg.OfTool.AnthropicContentFields) { + return true + } + // Check system messages (content is ContentUnion — []ChatCompletionContentPartTextParam or string). + if msg.OfSystem != nil { + if parts, ok := msg.OfSystem.Content.Value.([]openai.ChatCompletionContentPartTextParam); ok { + for j := range parts { + if isCacheEnabled(parts[j].AnthropicContentFields) { + return true + } + } + } + } + // Check user messages (content is StringOrUserRoleContentUnion — + // []ChatCompletionContentPartUserUnionParam or string). + if msg.OfUser != nil { + if parts, ok := msg.OfUser.Content.Value.([]openai.ChatCompletionContentPartUserUnionParam); ok { + for j := range parts { + part := &parts[j] + if part.OfText != nil && isCacheEnabled(part.OfText.AnthropicContentFields) { + return true + } + if part.OfImageURL != nil && isCacheEnabled(part.OfImageURL.AnthropicContentFields) { + return true + } + if part.OfInputAudio != nil && isCacheEnabled(part.OfInputAudio.AnthropicContentFields) { + return true + } + } + } + } + } + return false +} + // applyVendorSpecificFields applies GCP Vertex AI vendor-specific fields to the Gemini request. // These fields allow users to access advanced GCP-specific features not available in the OpenAI API. // Vendor fields override any conflicting fields that were set during the standard translation process. @@ -584,6 +632,9 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) applyVendorSpecificField if gcpVendorFields.SafetySettings != nil { gcr.SafetySettings = gcpVendorFields.SafetySettings } + if gcpVendorFields.CachedContent != "" { + gcr.CachedContent = gcpVendorFields.CachedContent + } } func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) geminiResponseToOpenAIMessage(gcr *genai.GenerateContentResponse, responseModel string) (*openai.ChatCompletionResponse, error) { diff --git a/internal/translator/openai_gcpvertexai_test.go b/internal/translator/openai_gcpvertexai_test.go index 439c82d219..4b239c8ce3 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -12,6 +12,8 @@ import ( "strings" "testing" + anthropic "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/shared/constant" corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" "github.com/google/go-cmp/cmp" @@ -780,6 +782,71 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) }, wantBody: wantBdyWithEnterpriseWebSearch, }, + { + name: "Request with explicit cachedContent vendor field", + input: openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + MaxTokens: ptr.To(int64(1024)), + Messages: []openai.ChatCompletionMessageParamUnion{ + { + OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{Value: "Summarize the document."}, + }, + }, + }, + GCPVertexAIVendorFields: &openai.GCPVertexAIVendorFields{ + CachedContent: "projects/my-project/locations/us-central1/cachedContents/abc123", + }, + }, + onRetry: false, + wantError: false, + wantHeaderMut: []internalapi.Header{ + {":path", "publishers/google/models/gemini-1.5-pro:generateContent"}, + }, + wantBody: []byte(`{ + "contents": [{"parts": [{"text": "Summarize the document."}], "role": "user"}], + "tools": null, + "generationConfig": {"maxOutputTokens": 1024}, + "cachedContent": "projects/my-project/locations/us-central1/cachedContents/abc123" + }`), + }, + { + name: "Request with both cache_control markers and explicit cachedContent returns error", + input: openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + MaxTokens: ptr.To(int64(1024)), + Messages: []openai.ChatCompletionMessageParamUnion{ + { + OfSystem: &openai.ChatCompletionSystemMessageParam{ + Role: openai.ChatMessageRoleSystem, + Content: openai.ContentUnion{Value: []openai.ChatCompletionContentPartTextParam{ + { + Type: "text", + Text: "You are a helpful assistant.", + AnthropicContentFields: &openai.AnthropicContentFields{ + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + }, + }, + }}, + }, + }, + { + OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{Value: "Hello"}, + }, + }, + }, + GCPVertexAIVendorFields: &openai.GCPVertexAIVendorFields{ + CachedContent: "projects/my-project/locations/us-central1/cachedContents/abc123", + }, + }, + onRetry: true, + wantError: true, + }, } for _, tc := range tests { @@ -819,6 +886,118 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) } } +func TestGCPRequestHasCacheControlMarkers(t *testing.T) { + ephemeral := anthropic.CacheControlEphemeralParam{Type: constant.ValueOf[constant.Ephemeral]()} + + tests := []struct { + name string + req openai.ChatCompletionRequest + want bool + }{ + { + name: "no cache_control markers", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{Value: "hello"}, + }}, + }, + }, + want: false, + }, + { + name: "cache_control on system message content part", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfSystem: &openai.ChatCompletionSystemMessageParam{ + Role: openai.ChatMessageRoleSystem, + Content: openai.ContentUnion{Value: []openai.ChatCompletionContentPartTextParam{ + { + Type: "text", + Text: "You are helpful.", + AnthropicContentFields: &openai.AnthropicContentFields{ + CacheControl: ephemeral, + }, + }, + }}, + }}, + }, + }, + want: true, + }, + { + name: "cache_control on user message text content part", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{ + Value: []openai.ChatCompletionContentPartUserUnionParam{ + {OfText: &openai.ChatCompletionContentPartTextParam{ + Type: "text", + Text: "large context", + AnthropicContentFields: &openai.AnthropicContentFields{ + CacheControl: ephemeral, + }, + }}, + }, + }, + }}, + }, + }, + want: true, + }, + { + name: "cache_control on tool message", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfTool: &openai.ChatCompletionToolMessageParam{ + Role: openai.ChatMessageRoleTool, + ToolCallID: "call_1", + Content: openai.ContentUnion{Value: "result"}, + AnthropicContentFields: &openai.AnthropicContentFields{ + CacheControl: ephemeral, + }, + }}, + }, + }, + want: true, + }, + { + name: "string user message has no markers", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfUser: &openai.ChatCompletionUserMessageParam{ + Role: openai.ChatMessageRoleUser, + Content: openai.StringOrUserRoleContentUnion{Value: "plain string"}, + }}, + }, + }, + want: false, + }, + { + name: "system string content has no markers", + req: openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessageParamUnion{ + {OfSystem: &openai.ChatCompletionSystemMessageParam{ + Role: openai.ChatMessageRoleSystem, + Content: openai.ContentUnion{Value: "plain system string"}, + }}, + }, + }, + want: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := gcpRequestHasCacheControlMarkers(&tc.req) + require.Equal(t, tc.want, got) + }) + } +} + func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_ResponseHeaders(t *testing.T) { tests := []struct { name string From b3c075bb874d12efedbec4c42057427059857478 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Fri, 31 Jul 2026 13:58:03 +0530 Subject: [PATCH 02/17] gcpcache: add in-process Vertex AI context-cache resolver Introduce the internal/gcpcache package and the GCPAuthHandler interface required to surface GCP credentials to the resolver without duplicating credential management. Changes: - internal/filterapi/runtime.go: add GCPAuthHandler interface extending BackendAuthHandler with GCPTokenSource(), GCPRegion(), GCPProject() methods so in-process components can authenticate to GCP APIs. - internal/backendauth/gcp.go: implement GCPAuthHandler on gcpHandler; add GCPTokenSource() (wraps static token in StaticTokenSource), GCPRegion() and GCPProject() accessors. - internal/translator/gemini_helper.go: export OpenAIMessagesToGeminiContents and OpenAIToolsToGeminiTools so gcpcache can build the cached prefix in Gemini wire format without duplicating conversion logic. - internal/gcpcache/gcpcache.go: new CacheResolver interface and resolver implementation with breakpoint detection, prefix split, deterministic SHA-256 cache key (stored as Google cachedContents displayName), in-memory TTL memo, Google cachedContents list/create client, and post-create re-list to converge duplicate-create races across replicas. - internal/gcpcache/gcpcache_test.go: unit tests covering findBreakpoint, extractTTL, computeCacheKey determinism, cache miss+create, cache hit from Google list, memo hit (no API calls), memo expiry/refetch, create failure, TTL default/override, and duplicate-create race convergence. Signed-off-by: Sivanantham Chinnaiyan --- internal/backendauth/gcp.go | 16 + internal/filterapi/runtime.go | 15 + internal/gcpcache/gcpcache.go | 520 ++++++++++++++++++++++++ internal/gcpcache/gcpcache_test.go | 573 +++++++++++++++++++++++++++ internal/translator/gemini_helper.go | 13 + 5 files changed, 1137 insertions(+) create mode 100644 internal/gcpcache/gcpcache.go create mode 100644 internal/gcpcache/gcpcache_test.go diff --git a/internal/backendauth/gcp.go b/internal/backendauth/gcp.go index 4dd77c78f1..f9526c3578 100644 --- a/internal/backendauth/gcp.go +++ b/internal/backendauth/gcp.go @@ -96,3 +96,19 @@ func (g *gcpHandler) Do(_ context.Context, requestHeaders map[string]string, _ [ requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", accessToken) return []internalapi.Header{{":path", newPath}, {"Authorization", fmt.Sprintf("Bearer %s", accessToken)}}, nil } + +// GCPTokenSource implements [filterapi.GCPAuthHandler]. +// It returns the OAuth2 token source used to authenticate to GCP APIs. +// When a static access token is configured, it wraps it in a StaticTokenSource. +func (g *gcpHandler) GCPTokenSource() oauth2.TokenSource { + if g.tokenSource != nil { + return g.tokenSource + } + return oauth2.StaticTokenSource(&oauth2.Token{AccessToken: g.gcpAccessToken}) +} + +// GCPRegion implements [filterapi.GCPAuthHandler]. +func (g *gcpHandler) GCPRegion() string { return g.region } + +// GCPProject implements [filterapi.GCPAuthHandler]. +func (g *gcpHandler) GCPProject() string { return g.projectName } diff --git a/internal/filterapi/runtime.go b/internal/filterapi/runtime.go index 3885fb71f9..6958431056 100644 --- a/internal/filterapi/runtime.go +++ b/internal/filterapi/runtime.go @@ -10,6 +10,7 @@ import ( "fmt" "github.com/google/cel-go/cel" + "golang.org/x/oauth2" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/llmcostcel" @@ -22,6 +23,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) diff --git a/internal/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go new file mode 100644 index 0000000000..9c4e256b1d --- /dev/null +++ b/internal/gcpcache/gcpcache.go @@ -0,0 +1,520 @@ +// 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 in-process 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 up an in-memory TTL memo. +// 4. On a memo miss, lists Google cachedContents for the target region and model. +// 5. Creates the cache entry if not found, then returns the cache resource name and the +// non-cached messages. +// +// The CacheResolver interface allows the implementation to be replaced with an external +// cache service in the future without changing the request-path callers. +package gcpcache + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "net/http" + "sync" + "time" + + "google.golang.org/genai" + + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "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 default cache TTL when none is specified in the cache_control marker. + defaultTTL = "300s" + + // gcpCachedContentsBasePath is the base URL for the Vertex AI cachedContents REST API. + gcpCachedContentsBasePath = "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/cachedContents" +) + +// ResolveResult holds the outcome of a successful cache resolution. +type ResolveResult struct { + // CacheName is the full Google resource name of the resolved or created cache entry. + // Format: "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 Google's create response). + // Only populated when Created is true. + TokenCount int + // ExpireTime is the cache expiration time reported by Google. + ExpireTime time.Time +} + +// CacheResolver resolves or creates a GCP Vertex AI cached content entry for requests +// that carry Anthropic-style cache_control markers. +type CacheResolver interface { + // Resolve inspects openAIReq for cache_control markers, resolves or creates the + // corresponding Google cachedContents entry, and returns the result. + // The caller is responsible for injecting ResolveResult.CacheName as the + // cachedContent field on the Gemini request and replacing the request messages + // with ResolveResult.Messages. + Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) +} + +// memoEntry is an in-memory cache entry for resolved Google cache names. +type memoEntry struct { + cacheName string + expireTime time.Time +} + +// resolver is the default in-process CacheResolver implementation. +type resolver struct { + httpClient *http.Client + + mu sync.Mutex + memo map[string]memoEntry // key → memoEntry; guarded by mu +} + +// New creates a new CacheResolver. The provided httpClient is used for calls to the +// Google cachedContents API; pass nil to use the default http.Client. +func New(httpClient *http.Client) CacheResolver { + if httpClient == nil { + httpClient = &http.Client{Timeout: 30 * time.Second} + } + return &resolver{ + httpClient: httpClient, + memo: make(map[string]memoEntry), + } +} + +// Resolve implements CacheResolver. +func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) { + region := gcpAuth.GCPRegion() + project := gcpAuth.GCPProject() + + // 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 = messages[0:breakpoint+1], remainder = messages[breakpoint+1:]. + cachedMessages := openAIReq.Messages[:breakpoint+1] + remainderMessages := openAIReq.Messages[breakpoint+1:] + + // Extract the TTL from the breakpoint message (fall back to default). + ttl := extractTTL(openAIReq.Messages[breakpoint]) + + // Build the Gemini cached prefix (tools + system + contents). + contents, systemInstruction, err := translator.OpenAIMessagesToGeminiContents(cachedMessages, openAIReq.Model) + if err != nil { + return nil, fmt.Errorf("gcpcache: converting cached messages to Gemini format: %w", err) + } + geminiTools, err := translator.OpenAIToolsToGeminiTools(openAIReq.Tools, false) + if err != nil { + return nil, fmt.Errorf("gcpcache: converting tools to Gemini format: %w", err) + } + + // Compute a deterministic cache key. + cacheKey, err := computeCacheKey(openAIReq.Model, contents, systemInstruction, geminiTools) + if err != nil { + return nil, fmt.Errorf("gcpcache: computing cache key: %w", err) + } + + // Check in-memory memo first. + if entry, ok := r.getMemo(cacheKey); ok { + return &ResolveResult{ + CacheName: entry.cacheName, + Messages: remainderMessages, + ExpireTime: entry.expireTime, + }, nil + } + + // Get access token. + tokenSrc := gcpAuth.GCPTokenSource() + token, err := tokenSrc.Token() + if err != nil { + return nil, fmt.Errorf("gcpcache: getting GCP access token: %w", err) + } + accessToken := token.AccessToken + + // List existing caches and match by displayName (= cacheKey). + baseURL := fmt.Sprintf(gcpCachedContentsBasePath, region, project, region) + existingName, expireTime, err := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) + if err != nil { + return nil, fmt.Errorf("gcpcache: listing cached contents: %w", err) + } + + if existingName != "" { + r.setMemo(cacheKey, existingName, expireTime) + return &ResolveResult{ + CacheName: existingName, + Messages: remainderMessages, + ExpireTime: expireTime, + }, nil + } + + // Cache not found — create it. + 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: creating cached content: %w", err) + } + + // After create, re-list and prefer the oldest match to converge duplicate-create races. + finalName, finalExpire, listErr := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) + if listErr == nil && finalName != "" && finalName != created { + // Another replica created a cache with the same key; use the one from the list. + r.setMemo(cacheKey, finalName, finalExpire) + return &ResolveResult{ + CacheName: finalName, + Messages: remainderMessages, + ExpireTime: finalExpire, + }, nil + } + + r.setMemo(cacheKey, created, expireTime) + return &ResolveResult{ + CacheName: created, + Messages: remainderMessages, + Created: true, + TokenCount: tokenCount, + ExpireTime: expireTime, + }, nil +} + +// ----------------------------------------------------------------------- +// 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) 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)) + } + } + 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 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 + } +} + +// computeKeyInputs converts the cached message prefix into Gemini format for key generation. +// It is separated from Resolve so tests can reproduce the same key deterministically. +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 { + 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 +// (model, contents, systemInstruction, tools) tuple. The digest is used as the +// Google cachedContents displayName. +func computeCacheKey(model string, contents []genai.Content, systemInstruction *genai.Content, tools []genai.Tool) (string, error) { + input := cacheKeyInput{ + Model: model, + Contents: contents, + SystemInstruction: systemInstruction, + Tools: tools, + } + b, err := json.Marshal(input) + if err != nil { + return "", fmt.Errorf("marshaling cache key input: %w", err) + } + h := sha256.Sum256(b) + return hex.EncodeToString(h[:]), nil +} + +// ----------------------------------------------------------------------- +// In-memory memo +// ----------------------------------------------------------------------- + +func (r *resolver) getMemo(key string) (memoEntry, bool) { + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.memo[key] + if !ok { + return memoEntry{}, false + } + // Treat entries expiring within 10 s as stale so we have time to re-resolve. + if time.Until(entry.expireTime) < 10*time.Second { + delete(r.memo, key) + return memoEntry{}, false + } + return entry, true +} + +func (r *resolver) setMemo(key, cacheName string, expireTime time.Time) { + r.mu.Lock() + defer r.mu.Unlock() + r.memo[key] = memoEntry{cacheName: cacheName, expireTime: expireTime} +} + +// ----------------------------------------------------------------------- +// Google cachedContents REST API calls +// ----------------------------------------------------------------------- + +// cachedContentItem is the subset of the cachedContents list item we care about. +type cachedContentItem struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + ExpireTime string `json:"expireTime"` // RFC 3339 + Model string `json:"model"` +} + +type listResponse struct { + CachedContents []cachedContentItem `json:"cachedContents"` +} + +// listAndMatch lists cachedContents for the given region+project and returns the +// name of the first entry whose displayName matches cacheKey and whose model +// suffix matches the requested model. Returns ("", zero, nil) when not found. +func (r *resolver) listAndMatch(ctx context.Context, baseURL, accessToken, cacheKey, model string) (string, time.Time, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL, nil) + if err != nil { + return "", time.Time{}, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + + resp, err := r.httpClient.Do(req) + if err != nil { + return "", time.Time{}, err + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", time.Time{}, err + } + if resp.StatusCode != http.StatusOK { + return "", time.Time{}, fmt.Errorf("list cachedContents returned HTTP %d: %s", resp.StatusCode, body) + } + + var lr listResponse + if err = json.Unmarshal(body, &lr); err != nil { + return "", time.Time{}, fmt.Errorf("decoding list response: %w", err) + } + + for _, item := range lr.CachedContents { + if item.DisplayName != cacheKey { + continue + } + // Model in the item is the full resource name; check the suffix. + if !modelMatchesSuffix(item.Model, model) { + continue + } + expireTime, _ := time.Parse(time.RFC3339, item.ExpireTime) + return item.Name, expireTime, nil + } + return "", time.Time{}, nil +} + +// modelMatchesSuffix checks whether the full model resource name ends with the +// requested short model name (e.g. "publishers/google/models/gemini-1.5-pro"). +func modelMatchesSuffix(fullModel, shortModel string) bool { + if fullModel == shortModel { + return true + } + suffix := "models/" + shortModel + return len(fullModel) >= len(suffix) && fullModel[len(fullModel)-len(suffix):] == suffix +} + +// createCacheRequest is the request body for POST .../cachedContents. +type createCacheRequest struct { + Model string `json:"model"` + Contents []genai.Content `json:"contents"` + SystemInstruction *genai.Content `json:"systemInstruction,omitempty"` + Tools []genai.Tool `json:"tools,omitempty"` + DisplayName string `json:"displayName"` + TTL string `json:"ttl"` +} + +// createCacheResponse is the subset of the create response we read. +type createCacheResponse struct { + Name string `json:"name"` + ExpireTime string `json:"expireTime"` + UsageMetadata *cacheUsageMetadata `json:"usageMetadata,omitempty"` +} + +type cacheUsageMetadata struct { + TotalTokenCount int `json:"totalTokenCount"` +} + +// createCache calls POST .../cachedContents 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 := createCacheRequest{ + 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("marshaling create request: %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 createCacheResponse + if err = json.Unmarshal(respBody, &cr); err != nil { + return "", 0, time.Time{}, fmt.Errorf("decoding create response: %w", err) + } + + expireTime, _ = time.Parse(time.RFC3339, cr.ExpireTime) + if cr.UsageMetadata != nil { + tokenCount = cr.UsageMetadata.TotalTokenCount + } + return cr.Name, tokenCount, expireTime, nil +} diff --git a/internal/gcpcache/gcpcache_test.go b/internal/gcpcache/gcpcache_test.go new file mode 100644 index 0000000000..444fa76e2a --- /dev/null +++ b/internal/gcpcache/gcpcache_test.go @@ -0,0 +1,573 @@ +// 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" + "net/http" + "net/http/httptest" + "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" + + "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "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 + listCalls int + createCalls int +} + +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++ + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(f.listResponse)) + case http.MethodPost: + f.createCalls++ + 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) +} + +func resolverWithServer(srvURL string) *resolver { + host := srvURL[len("http://"):] + return New(&http.Client{ + Transport: &redirectTransport{fakeHost: host}, + Timeout: 5 * time.Second, + }).(*resolver) +} + +// 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(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)) + }) + } +} + +// ----------------------------------------------------------------------- +// 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("gemini-1.5-pro", contents, sys, nil) + require.NoError(t, err) + k2, err := computeCacheKey("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("gemini-1.5-pro", c, s, nil) + k2, _ := computeCacheKey("gemini-2.0-flash", c, s, nil) + assert.NotEqual(t, k1, k2) +} + +// ----------------------------------------------------------------------- +// Integration tests: Resolve() against fake HTTP server +// ----------------------------------------------------------------------- + +func TestResolver_NoMarkers_ReturnsNil(t *testing.T) { + r := New(nil).(*resolver) + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{userMsg("hello")}, + } + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + res, err := r.Resolve(context.Background(), req, auth) + require.NoError(t, err) + assert.Nil(t, res, "no cache_control markers should return nil") +} + +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) + + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + 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, auth) + 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) + // list → create → re-list (post-create convergence check) + assert.Equal(t, 2, fake.listCalls) + assert.Equal(t, 1, fake.createCalls) +} + +func TestResolver_CacheHit_FromGoogleList(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) + + fake := newFakeCacheServer(t, string(listBody), "", http.StatusOK) + r := resolverWithServer(fake.srv.URL) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + + res, err := r.Resolve(context.Background(), req, auth) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/existing", res.CacheName) + assert.False(t, res.Created) + assert.Equal(t, 0, fake.createCalls) +} + +func TestResolver_MemoHit_SkipsGoogleAPICalls(t *testing.T) { + fake := newFakeCacheServer(t, `{"cachedContents":[]}`, "", http.StatusOK) + r := resolverWithServer(fake.srv.URL) + + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("Cached system.", ephemeralFields()), + userMsg("q"), + }, + } + key := computeKeyForRequest(t, req) + r.setMemo(key, "projects/p/locations/r/cachedContents/memo-hit", time.Now().Add(10*time.Minute)) + + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + res, err := r.Resolve(context.Background(), req, auth) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, "projects/p/locations/r/cachedContents/memo-hit", res.CacheName) + assert.False(t, res.Created) + assert.Equal(t, 0, fake.listCalls, "memo hit must not call the Google API") + assert.Equal(t, 0, fake.createCalls) +} + +func TestResolver_MemoExpiry_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) + + listBody, _ := json.Marshal(map[string]interface{}{ + "cachedContents": []map[string]interface{}{ + { + "name": "projects/p/locations/us-central1/cachedContents/refreshed", + "displayName": key, + "model": "publishers/google/models/gemini-1.5-pro", + "expireTime": expireISO, + }, + }, + }) + fake := newFakeCacheServer(t, string(listBody), "", http.StatusOK) + r := resolverWithServer(fake.srv.URL) + // Seed with entry expiring within the 10 s stale window → should be evicted. + r.setMemo(key, "projects/p/locations/us-central1/cachedContents/stale", time.Now().Add(5*time.Second)) + + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + res, err := r.Resolve(context.Background(), req, auth) + 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.listCalls, "stale memo must trigger a list call") +} + +func TestResolver_CreateFailure_ReturnsError(t *testing.T) { + fake := newFakeCacheServer(t, + `{"cachedContents":[]}`, + `{"error":{"message":"below minimum tokens"}}`, + http.StatusUnprocessableEntity, + ) + r := resolverWithServer(fake.srv.URL) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + _, err := r.Resolve(context.Background(), req, auth) + 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) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + 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, auth) + require.NoError(t, err) + + var body createCacheRequest + require.NoError(t, json.Unmarshal(capturedBody, &body)) + assert.Equal(t, defaultTTL, body.TTL) +} + +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) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFieldsWithTTL("1h")), + userMsg("q"), + }, + } + _, err := r.Resolve(context.Background(), req, auth) + require.NoError(t, err) + + var body createCacheRequest + require.NoError(t, json.Unmarshal(capturedBody, &body)) + assert.Equal(t, "3600s", body.TTL) +} + +func TestResolver_DuplicateCreateRace_UsesListResult(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) + + listCalls := 0 + // First list: empty (cache not yet created by this replica). + // Create: returns "new-by-us". + // Second list (post-create): returns "older-by-other" — simulates another replica winning the race. + listBodyEmpty := `{"cachedContents":[]}` + 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++ + w.WriteHeader(http.StatusOK) + if listCalls == 1 { + _, _ = w.Write([]byte(listBodyEmpty)) + } else { + _, _ = 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 + `"}`)) + } + })) + defer srv.Close() + + r := resolverWithServer(srv.URL) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + res, err := r.Resolve(context.Background(), req, auth) + require.NoError(t, err) + require.NotNil(t, res) + // Resolver should prefer the entry from the post-create re-list (the other replica's entry). + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/older-by-other", res.CacheName) + // Created should be false when we converged to another replica's cache. + assert.False(t, res.Created) +} diff --git a/internal/translator/gemini_helper.go b/internal/translator/gemini_helper.go index 8ca1ffc550..fd6e65d5ae 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 From 25a0fe813d4f847950c632fc1b6a5bbed6c51c29 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Fri, 31 Jul 2026 16:45:47 +0530 Subject: [PATCH 03/17] extproc: wire GCP context-cache resolver into request/response path Add the GCPCacheSetter interface to the translator package so the GCP Vertex AI translator can receive a resolved cache entry before RequestBody is called. The translator injects the cachedContent resource name into the Gemini request, uses the filtered (non-cached) message remainder as the request contents, and records cache-write token counts in ResponseBody for both streaming and non-streaming responses. Populate a shared CacheResolver on RuntimeBackend (via Server.LoadConfig) for every GCP Vertex AI backend so the in-memory TTL memo is preserved across requests. In upstreamProcessor.ProcessRequestHeaders, after the translator is selected, attempt cache resolution when the resolver, a GCPAuthHandler, and a GCPCacheSetter-capable translator are all present. Resolver failures produce a 502 ImmediateResponse; nil results (no cache_control markers) are a no-op. Add unit tests Signed-off-by: Sivanantham Chinnaiyan --- internal/extproc/processor_impl.go | 37 ++- internal/extproc/processor_impl_test.go | 247 ++++++++++++++++++ internal/extproc/server.go | 10 + internal/filterapi/runtime.go | 5 + internal/translator/openai_gcpvertexai.go | 48 +++- .../translator/openai_gcpvertexai_test.go | 103 ++++++++ internal/translator/translator.go | 24 ++ 7 files changed, 472 insertions(+), 2 deletions(-) diff --git a/internal/extproc/processor_impl.go b/internal/extproc/processor_impl.go index 5bcdb2d269..81b7a5941d 100644 --- a/internal/extproc/processor_impl.go +++ b/internal/extproc/processor_impl.go @@ -24,10 +24,12 @@ 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/endpointspec" "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/gcpcache" "github.com/envoyproxy/ai-gateway/internal/headermutator" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/json" @@ -128,6 +130,9 @@ type ( backendName string routeName string handler filterapi.BackendAuthHandler + // cacheResolver is the GCP context-cache resolver for this backend. + // Non-nil only for GCP Vertex AI backends; set in SetBackend. + cacheResolver gcpcache.CacheResolver // cost is the cost of the request that is accumulated during the processing of the response. costs metrics.TokenUsage // metrics tracking. @@ -359,8 +364,35 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque // We force the body mutation in the following cases: // * The request is a retry request because the body mutation might have happened the previous iteration. // * 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. + // the token usage is calculated correctly without being pypassed. 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. + if u.cacheResolver != nil { + if cacheSetter, ok := u.translator.(translator.GCPCacheSetter); ok { + if gcpAuth, ok := u.handler.(filterapi.GCPAuthHandler); ok { + if req, ok := any(u.parent.originalRequestBody).(*openai.ChatCompletionRequest); ok { + result, resolveErr := u.cacheResolver.Resolve(ctx, req, gcpAuth) + if resolveErr != nil { + u.logger.Error("GCP 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 { + cacheSetter.SetGCPCacheResult(&translator.GCPCacheResult{ + CacheName: result.CacheName, + FilteredMessages: result.Messages, + Created: result.Created, + WriteTokenCount: uint32(result.TokenCount), //nolint:gosec + }) + } + } + } + } + } + newHeaders, newBody, err := u.translator.RequestBody(u.parent.originalRequestBodyRaw, u.parent.originalRequestBody, forceBodyMutation) if err != nil { if userFacingErr := internalapi.GetUserFacingError(err); userFacingErr != nil { @@ -705,6 +737,9 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) SetBackend(c u.backendName = backend.Backend.Name u.routeName = routeName u.handler = backend.Handler + if cr, ok := backend.CacheResolver.(gcpcache.CacheResolver); ok { + u.cacheResolver = cr + } 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 4482ca98e6..6d4ebf6dd2 100644 --- a/internal/extproc/processor_impl_test.go +++ b/internal/extproc/processor_impl_test.go @@ -23,6 +23,7 @@ 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" @@ -31,6 +32,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/bodymutator" "github.com/envoyproxy/ai-gateway/internal/endpointspec" "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/gcpcache" "github.com/envoyproxy/ai-gateway/internal/headermutator" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/json" @@ -38,6 +40,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/metrics" "github.com/envoyproxy/ai-gateway/internal/testing/testotel" "github.com/envoyproxy/ai-gateway/internal/tracing/tracingapi" + "github.com/envoyproxy/ai-gateway/internal/translator" ) func TestNewFactory(t *testing.T) { @@ -2480,3 +2483,247 @@ func mustCompileCEL(t *testing.T, expr string) cel.Program { require.NoError(t, err) return prog } + +// --------------------------------------------------------------------------- +// GCP context-cache resolver wiring tests +// --------------------------------------------------------------------------- + +// mockGCPCacheResolver is a controllable gcpcache.CacheResolver for testing. +type mockGCPCacheResolver struct { + retResult *gcpcache.ResolveResult + retErr error + called bool +} + +func (m *mockGCPCacheResolver) Resolve(_ context.Context, _ *openai.ChatCompletionRequest, _ filterapi.GCPAuthHandler) (*gcpcache.ResolveResult, error) { + m.called = true + return m.retResult, m.retErr +} + +// mockGCPCacheTranslator extends mockTranslator with GCPCacheSetter support. +type mockGCPCacheTranslator struct { + mockTranslator + setCacheResult *translator.GCPCacheResult +} + +func (m *mockGCPCacheTranslator) SetGCPCacheResult(result *translator.GCPCacheResult) { + 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_GCPCacheResolver_notCalled verifies that the resolver is not invoked +// when the translator does not implement GCPCacheSetter. +func Test_upstreamProcessor_GCPCacheResolver_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 := &mockGCPCacheResolver{} + // mockTranslator does not implement GCPCacheSetter. + 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 GCPCacheSetter") +} + +// Test_upstreamProcessor_GCPCacheResolver_noMarkers verifies that when the resolver returns nil +// (no cache_control markers), the translator is called without a cache result. +func Test_upstreamProcessor_GCPCacheResolver_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 := &mockGCPCacheResolver{retResult: nil} // no markers + tr := &mockGCPCacheTranslator{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, "SetGCPCacheResult must not be called when resolver returns nil") +} + +// Test_upstreamProcessor_GCPCacheResolver_hit verifies that a cache hit is injected into the translator. +func Test_upstreamProcessor_GCPCacheResolver_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 := &mockGCPCacheResolver{retResult: &gcpcache.ResolveResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/abc123", + Messages: filteredMessages, + Created: false, + }} + tr := &mockGCPCacheTranslator{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_GCPCacheResolver_created verifies that cache-write tokens are propagated. +func Test_upstreamProcessor_GCPCacheResolver_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 := &mockGCPCacheResolver{retResult: &gcpcache.ResolveResult{ + CacheName: "projects/p/locations/us-central1/cachedContents/new456", + Messages: []openai.ChatCompletionMessageParamUnion{}, + Created: true, + TokenCount: 512, + }} + tr := &mockGCPCacheTranslator{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) +} + +// Test_upstreamProcessor_GCPCacheResolver_error verifies that a resolver error produces a 502. +func Test_upstreamProcessor_GCPCacheResolver_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 := &mockGCPCacheResolver{retErr: errors.New("GCP API quota exceeded")} + tr := &mockGCPCacheTranslator{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) +} + +// Test_upstreamProcessor_GCPCacheResolver_noHandler verifies the resolver is skipped +// when the handler is not a GCPAuthHandler. +func Test_upstreamProcessor_GCPCacheResolver_noHandler(t *testing.T) { + someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) + var body openai.ChatCompletionRequest + require.NoError(t, json.Unmarshal(someBody, &body)) + + cr := &mockGCPCacheResolver{} + tr := &mockGCPCacheTranslator{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: &mockBackendAuthHandler{}, // not a GCPAuthHandler + cacheResolver: cr, + } + _, err := p.ProcessRequestHeaders(t.Context(), nil) + require.NoError(t, err) + require.False(t, cr.called, "resolver must not be called when handler is not GCPAuthHandler") + require.Nil(t, tr.setCacheResult) +} diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 1ef537cf2b..14b2a004be 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -27,6 +27,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/backendauth" "github.com/envoyproxy/ai-gateway/internal/filterapi" + "github.com/envoyproxy/ai-gateway/internal/gcpcache" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/redaction" ) @@ -83,6 +84,15 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error if err != nil { return fmt.Errorf("cannot create runtime filter config: %w", err) } + + // Attach a shared in-process context-cache resolver for each GCP Vertex AI backend. + // The resolver is reused across requests so the in-memory TTL memo is preserved. + for _, rb := range newConfig.Backends { + if _, ok := rb.Handler.(filterapi.GCPAuthHandler); ok { + rb.CacheResolver = gcpcache.New(nil) + } + } + s.config = newConfig // This is racey, but we don't care. return nil } diff --git a/internal/filterapi/runtime.go b/internal/filterapi/runtime.go index 6958431056..caa061e00b 100644 --- a/internal/filterapi/runtime.go +++ b/internal/filterapi/runtime.go @@ -68,6 +68,11 @@ type RuntimeBackend struct { Backend *Backend // Handler is the backend auth handler. Handler BackendAuthHandler + // CacheResolver is the in-process context-cache resolver for this backend. + // It is populated by the extproc layer for GCP Vertex AI backends (where Handler + // implements GCPAuthHandler). Held as any to avoid an import cycle between + // filterapi and gcpcache; callers type-assert to gcpcache.CacheResolver. + CacheResolver any } // RuntimeGlobalRequestCost is the configuration for gateway-level default request costs. diff --git a/internal/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index b1039031e4..37f0f10c52 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 + // pendingCacheResult is set by SetGCPCacheResult 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. + pendingCacheResult *GCPCacheResult + // cacheWriteTokens is the number of tokens stored in a newly created cache entry. + // Set during RequestBody (from pendingCacheResult) and consumed in ResponseBody to + // record cache-write cost via tokenUsage.SetCacheCreationInputTokens. + cacheWriteTokens uint32 } // RequestBody implements [OpenAIChatCompletionTranslator.RequestBody] for GCP Gemini. @@ -126,6 +134,14 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) RequestBody(_ []byte, op if err != nil { return nil, nil, err } + + // Capture cache-write token count for cost attribution in ResponseBody, then clear + // the pending result so retries start fresh (the resolver will re-check the memo). + if o.pendingCacheResult != nil { + o.cacheWriteTokens = o.pendingCacheResult.WriteTokenCount + o.pendingCacheResult = nil + } + newBody, err = json.Marshal(gcpReq) if err != nil { return nil, nil, fmt.Errorf("error marshaling Gemini request: %w", err) @@ -202,6 +218,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) @@ -266,6 +286,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) + } } } @@ -523,8 +547,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 is pending (set by SetGCPCacheResult), use the filtered (non-cached) + // messages so that only the new turns are sent in the Gemini request body. + messages := openAIReq.Messages + if o.pendingCacheResult != nil && o.pendingCacheResult.FilteredMessages != nil { + messages = o.pendingCacheResult.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 } @@ -565,6 +596,13 @@ 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 SetGCPCacheResult). + // This happens after vendor-field processing so an explicit vendor cachedContent still wins + // (the validation in RequestBody already rejected the mixed case). + if o.pendingCacheResult != nil && gcr.CachedContent == "" { + gcr.CachedContent = o.pendingCacheResult.CacheName + } + return &gcr, nil } @@ -637,6 +675,14 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) applyVendorSpecificField } } +// SetGCPCacheResult implements [GCPCacheSetter]. 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) SetGCPCacheResult(result *GCPCacheResult) { + o.pendingCacheResult = 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 4b239c8ce3..bb7b235308 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -2921,3 +2921,106 @@ func TestGCPVertexAIRedactBody(t *testing.T) { require.NotContains(t, *resp.Choices[0].Message.Content, "[REDACTED") }) } + +// TestOpenAIToGCPVertexAITranslator_GCPCacheSetter verifies that SetGCPCacheResult +// 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_GCPCacheSetter(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.SetGCPCacheResult(&GCPCacheResult{ + 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"`) + + // pendingCacheResult should be cleared after RequestBody. + require.Nil(t, tr.pendingCacheResult) + // cacheWriteTokens should be set for response attribution. + require.Equal(t, uint32(256), tr.cacheWriteTokens) + }) + + t.Run("cache hit does not set write tokens", func(t *testing.T) { + tr := NewChatCompletionOpenAIToGCPVertexAITranslator("").(*openAIToGCPVertexAITranslatorV1ChatCompletion) + + tr.SetGCPCacheResult(&GCPCacheResult{ + 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("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 5555c1f95c..4fcdc60571 100644 --- a/internal/translator/translator.go +++ b/internal/translator/translator.go @@ -117,6 +117,30 @@ type ResponseRedactor interface { RedactBody(resp *openai.ChatCompletionResponse) *openai.ChatCompletionResponse } +// GCPCacheSetter is an optional interface implemented by the GCP Vertex AI translator. +// The upstream processor calls SetGCPCacheResult after the context-cache resolver runs so +// that the translator can (a) inject the cache resource name into the Gemini request and +// (b) record cache-write token counts for cost attribution in the response phase. +type GCPCacheSetter interface { + // SetGCPCacheResult forwards the resolved (or newly created) cache entry to the translator. + // Must be called before RequestBody. + SetGCPCacheResult(result *GCPCacheResult) +} + +// GCPCacheResult carries the information returned by the cache resolver that the translator needs. +type GCPCacheResult struct { + // CacheName is the full Google resource name of 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 Gemini 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 Google'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 { From af29ee4f98332bdbdf264a6338e6fabdb5eaed64 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Fri, 31 Jul 2026 17:04:18 +0530 Subject: [PATCH 04/17] api/filterapi: add GCPContextCaching configuration surface Add a GCPContextCachingSpec to AIServiceBackend in both api/v1beta1 and api/v1alpha1. The field carries Enabled (bool) and DefaultTTL (a GCP seconds-format duration string validated by a kubebuilder regex pattern). The controller wires it into filterapi.Backend.GCPContextCaching so the extproc layer can gate cache-resolver creation on the feature flag. In extproc/server.go, the shared CacheResolver is now only attached when GCPContextCaching.Enabled is true, keeping the resolver nil (and all cachedContents API calls suppressed) for backends that do not opt in. Generated: run make apigen codegen apidoc to regenerate CRD YAML, deepcopy functions, typed clients/listers/informers, and API reference. Tests added: - tests/crdcel: two AIServiceBackend fixtures (valid and invalid defaultTTL) exercised against the live envtest apiserver to verify CRD validation. - internal/controller/gateway_test.go: two reconcileFilterConfigSecret tests asserting that GCPContextCaching is populated when set and nil when absent. - internal/extproc/server_test.go: LoadConfig test asserting resolver creation is gated on Enabled=true / absent field / non-GCP backend. Signed-off-by: Sivanantham Chinnaiyan --- api/v1alpha1/ai_service_backend.go | 26 ++++ api/v1alpha1/zz_generated.deepcopy.go | 20 +++ api/v1beta1/ai_service_backend.go | 25 ++++ api/v1beta1/zz_generated.deepcopy.go | 20 +++ internal/controller/gateway.go | 8 ++ internal/controller/gateway_test.go | 116 ++++++++++++++++++ internal/extproc/server.go | 9 +- internal/extproc/server_test.go | 55 +++++++++ internal/filterapi/filterconfig.go | 16 +++ ...teway.envoyproxy.io_aiservicebackends.yaml | 48 ++++++++ site/docs/api/api.mdx | 64 ++++++++++ tests/crdcel/main_test.go | 5 + .../gcp-context-caching-enabled.yaml | 21 ++++ .../gcp-context-caching-invalid-ttl.yaml | 21 ++++ 14 files changed, 451 insertions(+), 3 deletions(-) create mode 100644 tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml create mode 100644 tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml diff --git a/api/v1alpha1/ai_service_backend.go b/api/v1alpha1/ai_service_backend.go index cddcc253f9..a7d61d2cac 100644 --- a/api/v1alpha1/ai_service_backend.go +++ b/api/v1alpha1/ai_service_backend.go @@ -76,10 +76,36 @@ type AIServiceBackendSpec struct { // +optional BodyMutation *HTTPBodyMutation `json:"bodyMutation,omitempty"` + // GCPContextCaching configures Gemini context caching for this backend. + // When set, the ai-gateway will automatically resolve or create a Google cachedContents + // entry for requests that carry Anthropic-style cache_control markers, and inject the + // resolved cache resource name into the Gemini API request. + // This field is only meaningful for GCPVertexAI backends. + // +optional + GCPContextCaching *GCPContextCachingSpec `json:"gcpContextCaching,omitempty"` + // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } +// GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. +type GCPContextCachingSpec struct { + // Enabled controls whether context caching is active for this backend. + // When false, cache_control markers in requests are ignored and no cachedContents + // API calls are made. + // + // +kubebuilder:validation:Required + Enabled bool `json:"enabled"` + + // DefaultTTL is the default time-to-live for newly created cache entries, expressed + // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. + // Must be at least "60s" per GCP API requirements. + // + // +optional + // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` + DefaultTTL string `json:"defaultTTL,omitempty"` +} + // HTTPHeaderMutation defines the mutation of HTTP headers that will be applied to the request type HTTPHeaderMutation struct { // Set overwrites/adds the request with the given header (name, value) diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 3c44727050..07ed7aedc7 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -338,6 +338,11 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { *out = new(HTTPBodyMutation) (*in).DeepCopyInto(*out) } + if in.GCPContextCaching != nil { + in, out := &in.GCPContextCaching, &out.GCPContextCaching + *out = new(GCPContextCachingSpec) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AIServiceBackendSpec. @@ -706,6 +711,21 @@ func (in *BackendSecurityPolicyStatus) DeepCopy() *BackendSecurityPolicyStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *GCPContextCachingSpec) DeepCopyInto(out *GCPContextCachingSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new GCPContextCachingSpec. +func (in *GCPContextCachingSpec) DeepCopy() *GCPContextCachingSpec { + if in == nil { + return nil + } + out := new(GCPContextCachingSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *GCPCredentialsFile) DeepCopyInto(out *GCPCredentialsFile) { *out = *in diff --git a/api/v1beta1/ai_service_backend.go b/api/v1beta1/ai_service_backend.go index ed0d3b8573..9b3a4c5833 100644 --- a/api/v1beta1/ai_service_backend.go +++ b/api/v1beta1/ai_service_backend.go @@ -90,7 +90,32 @@ type AIServiceBackendSpec struct { // +listMapKey=name // +kubebuilder:validation:MaxItems=16 HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` + // GCPContextCaching configures Gemini context caching for this backend. + // When set, the ai-gateway will automatically resolve or create a Google cachedContents + // entry for requests that carry Anthropic-style cache_control markers, and inject the + // resolved cache resource name into the Gemini API request. + // This field is only meaningful for GCPVertexAI backends. + // +optional + GCPContextCaching *GCPContextCachingSpec `json:"gcpContextCaching,omitempty"` // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } + +// GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. +type GCPContextCachingSpec struct { + // Enabled controls whether context caching is active for this backend. + // When false, cache_control markers in requests are ignored and no cachedContents + // API calls are made. + // + // +kubebuilder:validation:Required + Enabled bool `json:"enabled"` + + // DefaultTTL is the default time-to-live for newly created cache entries, expressed + // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. + // Must be at least "60s" per GCP API requirements. + // + // +optional + // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` + DefaultTTL string `json:"defaultTTL,omitempty"` +} diff --git a/api/v1beta1/zz_generated.deepcopy.go b/api/v1beta1/zz_generated.deepcopy.go index 10823978cf..3eac1ae533 100644 --- a/api/v1beta1/zz_generated.deepcopy.go +++ b/api/v1beta1/zz_generated.deepcopy.go @@ -345,6 +345,11 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.GCPContextCaching != nil { + in, out := &in.GCPContextCaching, &out.GCPContextCaching + *out = new(GCPContextCachingSpec) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AIServiceBackendSpec. @@ -783,6 +788,21 @@ func (in *CredentialOverrideFromRequestHeaders) DeepCopy() *CredentialOverrideFr return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *GCPContextCachingSpec) DeepCopyInto(out *GCPContextCachingSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new GCPContextCachingSpec. +func (in *GCPContextCachingSpec) DeepCopy() *GCPContextCachingSpec { + if in == nil { + return nil + } + out := new(GCPContextCachingSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *GCPCredentialsFile) DeepCopyInto(out *GCPCredentialsFile) { *out = *in diff --git a/internal/controller/gateway.go b/internal/controller/gateway.go index 255d6abc69..329bcd122c 100644 --- a/internal/controller/gateway.go +++ b/internal/controller/gateway.go @@ -523,6 +523,14 @@ func (c *GatewayController) reconcileFilterConfigSecret( b.Schema = schemaToFilterAPI(backendObj.Spec.APISchema) b.HeaderValueFilters = headerValueFiltersToFilterAPI(backendObj.Spec.HeaderValueFilters) + + // Wire GCP context caching configuration when present on the backend object. + if backendObj.Spec.GCPContextCaching != nil { + b.GCPContextCaching = &filterapi.GCPContextCaching{ + Enabled: backendObj.Spec.GCPContextCaching.Enabled, + DefaultTTL: backendObj.Spec.GCPContextCaching.DefaultTTL, + } + } } if bsp != nil { diff --git a/internal/controller/gateway_test.go b/internal/controller/gateway_test.go index ebe381462c..974f87d20a 100644 --- a/internal/controller/gateway_test.go +++ b/internal/controller/gateway_test.go @@ -4128,3 +4128,119 @@ func TestGatewayController_warnUndeclaredMetadataNamespaces(t *testing.T) { c.warnUndeclaredMetadataNamespaces(ec, []string{"declared.ns", "missing.ns", "other-missing.ns"}, "gw", "ns") require.Empty(t, logged) } + +// TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching verifies that the +// GCPContextCaching field on an AIServiceBackend is propagated to filterapi.Backend in the +// rendered filter configuration secret. +func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching(t *testing.T) { + const gwNamespace, someNamespace = "ns", "some-namespace" + fakeClient := requireNewFakeClientWithIndexes(t) + kube := fake2.NewClientset() + ctrl.SetLogger(zap.New(zap.UseFlagOptions(&zap.Options{Development: true, Level: zapcore.DebugLevel}))) + c := newTestGatewayController(fakeClient, kube, ctrl.Log, "envoy-gateway-system", + "docker.io/envoyproxy/ai-gateway-extproc:latest", "info", false, nil, true) + + routes := []aigv1b1.AIGatewayRoute{ + { + ObjectMeta: metav1.ObjectMeta{Name: "gcproute", Namespace: gwNamespace}, + Spec: aigv1b1.AIGatewayRouteSpec{ + Rules: []aigv1b1.AIGatewayRouteRule{ + { + BackendRefs: []aigv1b1.AIGatewayRouteRuleBackendRef{{Name: "gcp-backend"}}, + Matches: []aigv1b1.AIGatewayRouteRuleMatch{{}}, + }, + }, + }, + }, + } + + // Create the AIServiceBackend with GCPContextCaching enabled. + err := fakeClient.Create(t.Context(), &aigv1b1.AIServiceBackend{ + ObjectMeta: metav1.ObjectMeta{Name: "gcp-backend", Namespace: gwNamespace}, + Spec: aigv1b1.AIServiceBackendSpec{ + APISchema: aigv1b1.VersionedAPISchema{Name: aigv1b1.APISchemaGCPVertexAI}, + BackendRef: gwapiv1.BackendObjectReference{ + Name: "gcp-service", + Kind: ptr.To(gwapiv1.Kind("Backend")), + Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), + }, + GCPContextCaching: &aigv1b1.GCPContextCachingSpec{ + Enabled: true, + DefaultTTL: "600s", + }, + }, + }) + require.NoError(t, err) + + _, err = c.reconcileFilterConfigSecret(t.Context(), "gw", gwNamespace, someNamespace, routes, nil, "uuid1", nil, nil) + require.NoError(t, err) + + fc := requireFilterConfigFromBundle(t, kube, someNamespace, "gw", gwNamespace) + + // Find the backend that contains "gcp-backend" in its name. + var gcpBackend *filterapi.Backend + for i := range fc.Backends { + if strings.Contains(fc.Backends[i].Name, "gcp-backend") { + gcpBackend = &fc.Backends[i] + break + } + } + require.NotNil(t, gcpBackend, "expected to find gcp-backend in filter config") + require.NotNil(t, gcpBackend.GCPContextCaching, "GCPContextCaching must be populated in filter config") + require.True(t, gcpBackend.GCPContextCaching.Enabled) + require.Equal(t, "600s", gcpBackend.GCPContextCaching.DefaultTTL) +} + +// TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhenAbsent verifies that +// when GCPContextCaching is not set on the AIServiceBackend, filterapi.Backend.GCPContextCaching is nil. +func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhenAbsent(t *testing.T) { + const gwNamespace, someNamespace = "ns2", "some-namespace2" + fakeClient := requireNewFakeClientWithIndexes(t) + kube := fake2.NewClientset() + ctrl.SetLogger(zap.New(zap.UseFlagOptions(&zap.Options{Development: true, Level: zapcore.DebugLevel}))) + c := newTestGatewayController(fakeClient, kube, ctrl.Log, "envoy-gateway-system", + "docker.io/envoyproxy/ai-gateway-extproc:latest", "info", false, nil, true) + + routes := []aigv1b1.AIGatewayRoute{ + { + ObjectMeta: metav1.ObjectMeta{Name: "gcproute2", Namespace: gwNamespace}, + Spec: aigv1b1.AIGatewayRouteSpec{ + Rules: []aigv1b1.AIGatewayRouteRule{ + { + BackendRefs: []aigv1b1.AIGatewayRouteRuleBackendRef{{Name: "gcp-backend2"}}, + Matches: []aigv1b1.AIGatewayRouteRuleMatch{{}}, + }, + }, + }, + }, + } + + err := fakeClient.Create(t.Context(), &aigv1b1.AIServiceBackend{ + ObjectMeta: metav1.ObjectMeta{Name: "gcp-backend2", Namespace: gwNamespace}, + Spec: aigv1b1.AIServiceBackendSpec{ + APISchema: aigv1b1.VersionedAPISchema{Name: aigv1b1.APISchemaGCPVertexAI}, + BackendRef: gwapiv1.BackendObjectReference{ + Name: "gcp-service", + Kind: ptr.To(gwapiv1.Kind("Backend")), + Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), + }, + // GCPContextCaching intentionally omitted. + }, + }) + require.NoError(t, err) + + _, err = c.reconcileFilterConfigSecret(t.Context(), "gw", gwNamespace, someNamespace, routes, nil, "uuid2", nil, nil) + require.NoError(t, err) + + fc := requireFilterConfigFromBundle(t, kube, someNamespace, "gw", gwNamespace) + + var gcpBackend *filterapi.Backend + for i := range fc.Backends { + if strings.Contains(fc.Backends[i].Name, "gcp-backend2") { + gcpBackend = &fc.Backends[i] + break + } + } + require.NotNil(t, gcpBackend, "expected to find gcp-backend2 in filter config") + require.Nil(t, gcpBackend.GCPContextCaching, "GCPContextCaching must be nil when not set on the CRD") +} diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 14b2a004be..be4fcaf5e9 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -85,11 +85,14 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error return fmt.Errorf("cannot create runtime filter config: %w", err) } - // Attach a shared in-process context-cache resolver for each GCP Vertex AI backend. - // The resolver is reused across requests so the in-memory TTL memo is preserved. + // Attach a shared in-process context-cache resolver for each GCP Vertex AI backend + // that has context caching explicitly enabled. The resolver is reused across requests + // so the in-memory TTL memo is preserved. for _, rb := range newConfig.Backends { if _, ok := rb.Handler.(filterapi.GCPAuthHandler); ok { - rb.CacheResolver = gcpcache.New(nil) + if rb.Backend.GCPContextCaching != nil && rb.Backend.GCPContextCaching.Enabled { + rb.CacheResolver = gcpcache.New(nil) + } } } diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index 8ae104251a..da6cf04314 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -53,6 +53,61 @@ func TestServer_LoadConfig(t *testing.T) { require.NotNil(t, s.config) } +// TestServer_LoadConfig_GCPCacheResolver verifies that a CacheResolver is attached +// to a GCP backend only when GCPContextCaching.Enabled is true. +func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { + s := &Server{} + + config := &filterapi.Config{ + Backends: []filterapi.Backend{ + { + Name: "gcp-caching-on", + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{ + GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, + }, + GCPContextCaching: &filterapi.GCPContextCaching{Enabled: true, DefaultTTL: "600s"}, + }, + { + Name: "gcp-caching-off", + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{ + GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, + }, + GCPContextCaching: &filterapi.GCPContextCaching{Enabled: false}, + }, + { + Name: "gcp-no-caching-field", + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{ + GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, + }, + // GCPContextCaching is nil. + }, + { + Name: "non-gcp", + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}, + // No auth handler → not a GCPAuthHandler. + }, + }, + } + + err := s.LoadConfig(t.Context(), config) + require.NoError(t, err) + + rb := s.config.Backends["gcp-caching-on"] + require.NotNil(t, rb.CacheResolver, "CacheResolver must be set when GCPContextCaching.Enabled=true") + + rb = s.config.Backends["gcp-caching-off"] + require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when GCPContextCaching.Enabled=false") + + rb = s.config.Backends["gcp-no-caching-field"] + require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when GCPContextCaching is absent") + + rb = s.config.Backends["non-gcp"] + require.Nil(t, rb.CacheResolver, "CacheResolver must be nil for non-GCP backends") +} + func TestServer_Check(t *testing.T) { s, _ := requireNewServerWithMockProcessor(t) diff --git a/internal/filterapi/filterconfig.go b/internal/filterapi/filterconfig.go index f2b7775b6f..f6853c2712 100644 --- a/internal/filterapi/filterconfig.go +++ b/internal/filterapi/filterconfig.go @@ -207,6 +207,22 @@ type Backend struct { // HeaderValueFilters filter individual values out of multi-valued request headers before sending // the request to the backend. Optional. HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` + // GCPContextCaching configures Gemini context caching for this backend. Optional. + // When non-nil and Enabled is true, the extproc will resolve or create a Google + // cachedContents entry for requests that carry cache_control markers. + GCPContextCaching *GCPContextCaching `json:"gcpContextCaching,omitempty"` +} + +// GCPContextCaching configures in-process Gemini context caching for a GCP Vertex AI backend. +type GCPContextCaching struct { + // Enabled controls whether context caching is active for this backend. + // When false (or when this struct is nil), cache_control markers in requests + // are ignored and no cachedContents API calls are made. + Enabled bool `json:"enabled"` + // DefaultTTL is the default time-to-live for newly created cache entries expressed + // as a GCP duration string (e.g. "600s"). When empty, the resolver default (300s) is used. + // +optional + DefaultTTL string `json:"defaultTTL,omitempty"` } // BackendAuth corresponds partially to BackendSecurityPolicy in api/v1alpha1/api.go. diff --git a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml index 64a3fa17c1..52ac704620 100644 --- a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml +++ b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml @@ -224,6 +224,30 @@ spec: - path x-kubernetes-list-type: map type: object + gcpContextCaching: + description: |- + GCPContextCaching configures Gemini context caching for this backend. + When set, the ai-gateway will automatically resolve or create a Google cachedContents + entry for requests that carry Anthropic-style cache_control markers, and inject the + resolved cache resource name into the Gemini API request. + This field is only meaningful for GCPVertexAI backends. + properties: + defaultTTL: + description: |- + DefaultTTL is the default time-to-live for newly created cache entries, expressed + as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. + Must be at least "60s" per GCP API requirements. + pattern: ^[1-9][0-9]*s$ + type: string + enabled: + description: |- + Enabled controls whether context caching is active for this backend. + When false, cache_control markers in requests are ignored and no cachedContents + API calls are made. + type: boolean + required: + - enabled + type: object headerMutation: description: |- HeaderMutation defines the mutation of HTTP headers that will be applied to the request @@ -641,6 +665,30 @@ spec: - path x-kubernetes-list-type: map type: object + gcpContextCaching: + description: |- + GCPContextCaching configures Gemini context caching for this backend. + When set, the ai-gateway will automatically resolve or create a Google cachedContents + entry for requests that carry Anthropic-style cache_control markers, and inject the + resolved cache resource name into the Gemini API request. + This field is only meaningful for GCPVertexAI backends. + properties: + defaultTTL: + description: |- + DefaultTTL is the default time-to-live for newly created cache entries, expressed + as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. + Must be at least "60s" per GCP API requirements. + pattern: ^[1-9][0-9]*s$ + type: string + enabled: + description: |- + Enabled controls whether context caching is active for this backend. + When false, cache_control markers in requests are ignored and no cachedContents + API calls are made. + type: boolean + required: + - enabled + type: object headerMutation: description: |- HeaderMutation defines the mutation of HTTP headers that will be applied to the request diff --git a/site/docs/api/api.mdx b/site/docs/api/api.mdx index b559cebd54..47791b5711 100644 --- a/site/docs/api/api.mdx +++ b/site/docs/api/api.mdx @@ -579,6 +579,7 @@ QuotaPolicyList contains a list of QuotaPolicy - [BackendSecurityPolicySpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicyspec) - [BackendSecurityPolicyStatus](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicystatus) - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicytype) +- [GCPContextCachingSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpcontextcachingspec) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpcredentialsfile) - [GCPOIDCExchangeToken](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpoidcexchangetoken) - [GCPServiceAccountImpersonationConfig](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpserviceaccountimpersonationconfig) @@ -851,6 +852,11 @@ AIServiceBackendSpec details the AIServiceBackend configuration. type="[HTTPBodyMutation](#github-com-envoyproxy-ai-gateway-api-v1alpha1-httpbodymutation)" required="false" description="BodyMutation defines the mutation of HTTP request body JSON fields that will be applied to the request
before sending it to the backend." +/> @@ -1376,6 +1382,32 @@ BackendSecurityPolicyType specifies the type of auth mechanism used to access a required="false" description="" /> +#### GCPContextCachingSpec + + + +**Appears in:** +- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-aiservicebackendspec) + +GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. + +##### Fields + + + + + + #### GCPCredentialsFile @@ -3152,6 +3184,7 @@ MCPRouteList contains a list of MCPRoute. - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicytype) - [CredentialOverrideFromDynamicMetadata](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromdynamicmetadata) - [CredentialOverrideFromRequestHeaders](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromrequestheaders) +- [GCPContextCachingSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpcontextcachingspec) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpcredentialsfile) - [GCPOIDCExchangeToken](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpoidcexchangetoken) - [GCPServiceAccountImpersonationConfig](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpserviceaccountimpersonationconfig) @@ -3426,6 +3459,11 @@ AIServiceBackendSpec details the AIServiceBackend configuration. type="[HTTPHeaderValueFilter](#github-com-envoyproxy-ai-gateway-api-v1beta1-httpheadervaluefilter) array" required="false" description="HeaderValueFilters filter individual values out of multi-valued request headers before sending
the request to this backend, so that one value an upstream provider does not accept does not
fail the whole request. At most one filter per header name.
Only GCPAnthropic and AWSAnthropic backends honor this, and today only for the
`anthropic-beta` header — a filter naming any other header is accepted but has no effect.
Setting it on a backend with any other schema is rejected." +/> @@ -4049,6 +4087,32 @@ CredentialOverrideFromRequestHeaders sources the per-request credential from a r /> +#### GCPContextCachingSpec + + + +**Appears in:** +- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-aiservicebackendspec) + +GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. + +##### Fields + + + + + + #### GCPCredentialsFile diff --git a/tests/crdcel/main_test.go b/tests/crdcel/main_test.go index e238e04080..25a81cde72 100644 --- a/tests/crdcel/main_test.go +++ b/tests/crdcel/main_test.go @@ -95,6 +95,11 @@ func TestAIServiceBackends(t *testing.T) { }{ {name: "basic.yaml"}, {name: "anthropic-schema.yaml"}, + {name: "gcp-context-caching-enabled.yaml"}, + { + name: "gcp-context-caching-invalid-ttl.yaml", + expErr: "spec.gcpContextCaching.defaultTTL", + }, {name: "basic-eg-backend-aws.yaml"}, {name: "aws-openai-schema.yaml"}, {name: "basic-eg-backend-azure.yaml"}, diff --git a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml b/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml new file mode 100644 index 0000000000..745500215e --- /dev/null +++ b/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml @@ -0,0 +1,21 @@ +# 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. + +apiVersion: aigateway.envoyproxy.io/v1beta1 +kind: AIServiceBackend +metadata: + name: gcp-cache-backend + namespace: default +spec: + schema: + name: GCPVertexAI + backendRef: + name: gcp-service + kind: Backend + group: gateway.envoyproxy.io + port: 443 + gcpContextCaching: + enabled: true + defaultTTL: "600s" diff --git a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml b/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml new file mode 100644 index 0000000000..dad194c47c --- /dev/null +++ b/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml @@ -0,0 +1,21 @@ +# 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. + +apiVersion: aigateway.envoyproxy.io/v1beta1 +kind: AIServiceBackend +metadata: + name: gcp-cache-backend-bad-ttl + namespace: default +spec: + schema: + name: GCPVertexAI + backendRef: + name: gcp-service + kind: Backend + group: gateway.envoyproxy.io + port: 443 + gcpContextCaching: + enabled: true + defaultTTL: "not-a-duration" From a5c9315389652dee96ae093adf0a2e2f81944e18 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Tue, 4 Aug 2026 12:14:22 +0530 Subject: [PATCH 05/17] docs: update GCP context cache documentation URL Update the context cache documentation URL to point to the new docs.cloud.google.com location under gemini-enterprise-agent-platform. Signed-off-by: Sivanantham Chinnaiyan --- internal/apischema/gcp/gcp.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/apischema/gcp/gcp.go b/internal/apischema/gcp/gcp.go index d38ed0d9b0..3aaa21fd73 100644 --- a/internal/apischema/gcp/gcp.go +++ b/internal/apischema/gcp/gcp.go @@ -42,7 +42,7 @@ type GenerateContentRequest struct { // 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://cloud.google.com/vertex-ai/docs/context-cache/context-cache-overview + // https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/context-cache/context-cache-overview CachedContent string `json:"cachedContent,omitempty"` } From cb2e4d3f0306dff70066ec7aeb79a8d718d1cf28 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 5 Aug 2026 13:59:42 +0530 Subject: [PATCH 06/17] gcpcache: consolidate cached content types into gcp schema package Replace the local createCacheRequest/createCacheResponse/cacheUsageMetadata types in gcpcache with shared types defined in internal/apischema/gcp: - Add CreateCachedContent, CachedContent, CachedContentUsageMetadata, and EncryptionSpec to the gcp schema package so they can be reused across the codebase. - Set CreateCachedContent.TTL as string (GCP duration format e.g. "300s") rather than time.Duration, which marshals as a nanosecond integer and would be rejected by the Vertex AI REST API. - Set CachedContent.ExpireTime as time.Time so the response can be used directly without a separate time.Parse call. - Cast CachedContentUsageMetadata.TotalTokenCount (int32) to int when populating ResolveResult.TokenCount. - Update tests to unmarshal the POST body into gcp.CreateCachedContent instead of the now-removed local type. Signed-off-by: Sivanantham Chinnaiyan --- internal/apischema/gcp/gcp.go | 79 ++++++++++++++++++++++++++++++ internal/gcpcache/gcpcache.go | 62 ++++++++--------------- internal/gcpcache/gcpcache_test.go | 5 +- 3 files changed, 102 insertions(+), 44 deletions(-) diff --git a/internal/apischema/gcp/gcp.go b/internal/apischema/gcp/gcp.go index 3aaa21fd73..3f5c73027e 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" @@ -176,3 +178,80 @@ 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. + 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/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go index 9c4e256b1d..20b7f90c5b 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/gcpcache/gcpcache.go @@ -32,6 +32,7 @@ import ( "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/filterapi" "github.com/envoyproxy/ai-gateway/internal/json" @@ -111,27 +112,27 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion return nil, nil } - // Split: cached prefix = messages[0:breakpoint+1], remainder = messages[breakpoint+1:]. + // Split cached prefix. cachedMessages := openAIReq.Messages[:breakpoint+1] remainderMessages := openAIReq.Messages[breakpoint+1:] - // Extract the TTL from the breakpoint message (fall back to default). + // Extract the TTL from the breakpoint message. If not found, fall back to default. ttl := extractTTL(openAIReq.Messages[breakpoint]) // Build the Gemini cached prefix (tools + system + contents). contents, systemInstruction, err := translator.OpenAIMessagesToGeminiContents(cachedMessages, openAIReq.Model) if err != nil { - return nil, fmt.Errorf("gcpcache: converting cached messages to Gemini format: %w", err) + 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: converting tools to Gemini format: %w", err) + return nil, fmt.Errorf("gcpcache: failed to convert tools to Gemini format: %w", err) } // Compute a deterministic cache key. cacheKey, err := computeCacheKey(openAIReq.Model, contents, systemInstruction, geminiTools) if err != nil { - return nil, fmt.Errorf("gcpcache: computing cache key: %w", err) + return nil, fmt.Errorf("gcpcache: failed to compute cache key: %w", err) } // Check in-memory memo first. @@ -147,15 +148,15 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion tokenSrc := gcpAuth.GCPTokenSource() token, err := tokenSrc.Token() if err != nil { - return nil, fmt.Errorf("gcpcache: getting GCP access token: %w", err) + return nil, fmt.Errorf("gcpcache: failed to get GCP access token: %w", err) } accessToken := token.AccessToken - // List existing caches and match by displayName (= cacheKey). + // List existing caches and match by displayName (cacheKey). baseURL := fmt.Sprintf(gcpCachedContentsBasePath, region, project, region) existingName, expireTime, err := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) if err != nil { - return nil, fmt.Errorf("gcpcache: listing cached contents: %w", err) + return nil, fmt.Errorf("gcpcache: failed to list cached contents: %w", err) } if existingName != "" { @@ -170,7 +171,7 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // Cache not found — create it. 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: creating cached content: %w", err) + return nil, fmt.Errorf("gcpcache: failed to create cached content: %w", err) } // After create, re-list and prefer the oldest match to converge duplicate-create races. @@ -308,8 +309,7 @@ func anthropicTTLToGCP(ttl string) string { } } -// computeKeyInputs converts the cached message prefix into Gemini format for key generation. -// It is separated from Resolve so tests can reproduce the same key deterministically. +// 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 { @@ -341,7 +341,7 @@ func computeCacheKey(model string, contents []genai.Content, systemInstruction * } b, err := json.Marshal(input) if err != nil { - return "", fmt.Errorf("marshaling cache key input: %w", err) + return "", fmt.Errorf("failed to marshal cache key input: %w", err) } h := sha256.Sum256(b) return hex.EncodeToString(h[:]), nil @@ -413,7 +413,7 @@ func (r *resolver) listAndMatch(ctx context.Context, baseURL, accessToken, cache var lr listResponse if err = json.Unmarshal(body, &lr); err != nil { - return "", time.Time{}, fmt.Errorf("decoding list response: %w", err) + return "", time.Time{}, fmt.Errorf("failed to decode list response: %w", err) } for _, item := range lr.CachedContents { @@ -440,28 +440,7 @@ func modelMatchesSuffix(fullModel, shortModel string) bool { return len(fullModel) >= len(suffix) && fullModel[len(fullModel)-len(suffix):] == suffix } -// createCacheRequest is the request body for POST .../cachedContents. -type createCacheRequest struct { - Model string `json:"model"` - Contents []genai.Content `json:"contents"` - SystemInstruction *genai.Content `json:"systemInstruction,omitempty"` - Tools []genai.Tool `json:"tools,omitempty"` - DisplayName string `json:"displayName"` - TTL string `json:"ttl"` -} - -// createCacheResponse is the subset of the create response we read. -type createCacheResponse struct { - Name string `json:"name"` - ExpireTime string `json:"expireTime"` - UsageMetadata *cacheUsageMetadata `json:"usageMetadata,omitempty"` -} - -type cacheUsageMetadata struct { - TotalTokenCount int `json:"totalTokenCount"` -} - -// createCache calls POST .../cachedContents and returns the new cache's resource name, +// 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, @@ -474,7 +453,7 @@ func (r *resolver) createCache( // Vertex AI expects the full model resource name. fullModel := fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s", project, region, model) - body := createCacheRequest{ + body := gcp.CreateCachedContent{ Model: fullModel, Contents: contents, SystemInstruction: systemInstruction, @@ -484,7 +463,7 @@ func (r *resolver) createCache( } bodyBytes, err := json.Marshal(body) if err != nil { - return "", 0, time.Time{}, fmt.Errorf("marshaling create request: %w", err) + 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)) @@ -507,14 +486,13 @@ func (r *resolver) createCache( return "", 0, time.Time{}, fmt.Errorf("create cachedContent returned HTTP %d: %s", resp.StatusCode, respBody) } - var cr createCacheResponse + var cr gcp.CachedContent if err = json.Unmarshal(respBody, &cr); err != nil { - return "", 0, time.Time{}, fmt.Errorf("decoding create response: %w", err) + return "", 0, time.Time{}, fmt.Errorf("failed to decode create cache response: %w", err) } - expireTime, _ = time.Parse(time.RFC3339, cr.ExpireTime) if cr.UsageMetadata != nil { - tokenCount = cr.UsageMetadata.TotalTokenCount + tokenCount = int(cr.UsageMetadata.TotalTokenCount) } - return cr.Name, tokenCount, expireTime, nil + return cr.Name, tokenCount, cr.ExpireTime, nil } diff --git a/internal/gcpcache/gcpcache_test.go b/internal/gcpcache/gcpcache_test.go index 444fa76e2a..828d719809 100644 --- a/internal/gcpcache/gcpcache_test.go +++ b/internal/gcpcache/gcpcache_test.go @@ -18,6 +18,7 @@ import ( "github.com/stretchr/testify/require" "golang.org/x/oauth2" + "github.com/envoyproxy/ai-gateway/internal/apischema/gcp" "github.com/envoyproxy/ai-gateway/internal/apischema/openai" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/internalapi" @@ -475,7 +476,7 @@ func TestResolver_TTLDefault_SentToGoogle(t *testing.T) { _, err := r.Resolve(context.Background(), req, auth) require.NoError(t, err) - var body createCacheRequest + var body gcp.CreateCachedContent require.NoError(t, json.Unmarshal(capturedBody, &body)) assert.Equal(t, defaultTTL, body.TTL) } @@ -511,7 +512,7 @@ func TestResolver_TTLOverride_SentToGoogle(t *testing.T) { _, err := r.Resolve(context.Background(), req, auth) require.NoError(t, err) - var body createCacheRequest + var body gcp.CreateCachedContent require.NoError(t, json.Unmarshal(capturedBody, &body)) assert.Equal(t, "3600s", body.TTL) } From 2cebce3ad4c846272e7944cb45df670ed0c8a0cc Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Mon, 28 Sep 2026 13:09:43 +0530 Subject: [PATCH 07/17] Back GCP context cache with a shared Redis store The context cache previously kept its resolved cachedContents names in a per-replica in-memory map, so two gateway replicas seeing the same prefix each created their own Google cache. Replace the map with a CacheStore interface and a Redis-backed implementation that all replicas share. This changes the AIServiceBackend API in both v1alpha1 and v1beta1: gcpContextCaching is renamed to contextCache, its enabled field is removed, and a required url field gives the Redis address, either as host:port or as a redis:// URL. Setting contextCache now enables caching; leaving it unset disables it and cache_control markers are ignored. The filterapi config and controller wiring follow the same shape. Creation is arbitrated with SET NX PX: the replica that claims a key creates the cache while the others poll for its result, so a cold prefix arriving at N replicas produces one create rather than N. A waiter that times out resolves against Google itself, which degrades a stuck leader into an extra list rather than a duplicate create. Store failures are never fatal to a request. An unreachable Redis, or a url that does not parse, is treated as a cache miss and the request is served uncached, since caching is an optimization. Google-side failures keep failing fast, 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. Resolvers are now carried across config reloads, keyed by backend name and Redis URL, so a filterapi update does not tear down and rebuild each backend's connection pool. Adds github.com/redis/go-redis/v9, and github.com/alicebob/miniredis/v2 for tests. Signed-off-by: Sivanantham Chinnaiyan --- api/v1alpha1/ai_service_backend.go | 39 ++- api/v1alpha1/zz_generated.deepcopy.go | 14 +- api/v1beta1/ai_service_backend.go | 40 ++- api/v1beta1/zz_generated.deepcopy.go | 36 +- go.mod | 4 + go.sum | 12 + internal/controller/gateway.go | 10 +- internal/controller/gateway_test.go | 29 +- internal/extproc/server.go | 47 ++- internal/extproc/server_test.go | 90 ++++- internal/filterapi/filterconfig.go | 25 +- internal/gcpcache/gcpcache.go | 244 ++++++++++---- internal/gcpcache/gcpcache_test.go | 101 ++++-- internal/gcpcache/redis.go | 190 +++++++++++ internal/gcpcache/redis_test.go | 312 ++++++++++++++++++ internal/gcpcache/store.go | 48 +++ ...teway.envoyproxy.io_aiservicebackends.yaml | 58 ++-- site/docs/api/api.mdx | 90 ++--- 18 files changed, 1150 insertions(+), 239 deletions(-) create mode 100644 internal/gcpcache/redis.go create mode 100644 internal/gcpcache/redis_test.go create mode 100644 internal/gcpcache/store.go diff --git a/api/v1alpha1/ai_service_backend.go b/api/v1alpha1/ai_service_backend.go index a7d61d2cac..0178f59d55 100644 --- a/api/v1alpha1/ai_service_backend.go +++ b/api/v1alpha1/ai_service_backend.go @@ -76,27 +76,24 @@ type AIServiceBackendSpec struct { // +optional BodyMutation *HTTPBodyMutation `json:"bodyMutation,omitempty"` - // GCPContextCaching configures Gemini context caching for this backend. - // When set, the ai-gateway will automatically resolve or create a Google cachedContents + // ContextCache configures context caching for this backend. + // When set, the ai-gateway will automatically resolve or create a context cache // entry for requests that carry Anthropic-style cache_control markers, and inject the - // resolved cache resource name into the Gemini API request. - // This field is only meaningful for GCPVertexAI backends. + // resolved cache resource name into the API request. + // Currently only supported for GCP Vertex AI backends. // +optional - GCPContextCaching *GCPContextCachingSpec `json:"gcpContextCaching,omitempty"` + ContextCache *ContextCacheSpec `json:"contextCache,omitempty"` // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } -// GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. -type GCPContextCachingSpec struct { - // Enabled controls whether context caching is active for this backend. - // When false, cache_control markers in requests are ignored and no cachedContents - // API calls are made. - // - // +kubebuilder:validation:Required - Enabled bool `json:"enabled"` - +// ContextCacheSpec configures context caching for a backend. +// +// Setting this field enables context caching: cache_control markers in requests are +// resolved against the shared store at URL, and a cache entry is created when absent. +// Leaving it unset disables caching entirely, and markers are ignored. +type ContextCacheSpec struct { // DefaultTTL is the default time-to-live for newly created cache entries, expressed // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. // Must be at least "60s" per GCP API requirements. @@ -104,6 +101,20 @@ type GCPContextCachingSpec struct { // +optional // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` DefaultTTL string `json:"defaultTTL,omitempty"` + + // URL locates the shared Redis store that records which cachedContents entry a given + // request prefix resolved to. Sharing that state across gateway replicas is what + // prevents two replicas from independently creating a cache for the same prefix. + // + // It is given either as a bare "host:port" + // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. + // + // If Redis is unreachable at request time, caching is skipped and the request is + // served uncached rather than failed. + // + // +kubebuilder:validation:Required + // +kubebuilder:validation:MinLength=1 + URL string `json:"url"` } // HTTPHeaderMutation defines the mutation of HTTP headers that will be applied to the request diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 07ed7aedc7..280a9d9770 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -338,9 +338,9 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { *out = new(HTTPBodyMutation) (*in).DeepCopyInto(*out) } - if in.GCPContextCaching != nil { - in, out := &in.GCPContextCaching, &out.GCPContextCaching - *out = new(GCPContextCachingSpec) + if in.ContextCache != nil { + in, out := &in.ContextCache, &out.ContextCache + *out = new(ContextCacheSpec) **out = **in } } @@ -712,16 +712,16 @@ func (in *BackendSecurityPolicyStatus) DeepCopy() *BackendSecurityPolicyStatus { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *GCPContextCachingSpec) DeepCopyInto(out *GCPContextCachingSpec) { +func (in *ContextCacheSpec) DeepCopyInto(out *ContextCacheSpec) { *out = *in } -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new GCPContextCachingSpec. -func (in *GCPContextCachingSpec) DeepCopy() *GCPContextCachingSpec { +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ContextCacheSpec. +func (in *ContextCacheSpec) DeepCopy() *ContextCacheSpec { if in == nil { return nil } - out := new(GCPContextCachingSpec) + out := new(ContextCacheSpec) in.DeepCopyInto(out) return out } diff --git a/api/v1beta1/ai_service_backend.go b/api/v1beta1/ai_service_backend.go index 9b3a4c5833..46c20b3501 100644 --- a/api/v1beta1/ai_service_backend.go +++ b/api/v1beta1/ai_service_backend.go @@ -90,27 +90,25 @@ type AIServiceBackendSpec struct { // +listMapKey=name // +kubebuilder:validation:MaxItems=16 HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` - // GCPContextCaching configures Gemini context caching for this backend. - // When set, the ai-gateway will automatically resolve or create a Google cachedContents + + // ContextCache configures context caching for this backend. + // When set, the ai-gateway will automatically resolve or create a context cache // entry for requests that carry Anthropic-style cache_control markers, and inject the - // resolved cache resource name into the Gemini API request. - // This field is only meaningful for GCPVertexAI backends. + // resolved cache resource name into the API request. + // Currently only supported for GCP Vertex AI backends. // +optional - GCPContextCaching *GCPContextCachingSpec `json:"gcpContextCaching,omitempty"` + ContextCache *ContextCacheSpec `json:"contextCache,omitempty"` // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } -// GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. -type GCPContextCachingSpec struct { - // Enabled controls whether context caching is active for this backend. - // When false, cache_control markers in requests are ignored and no cachedContents - // API calls are made. - // - // +kubebuilder:validation:Required - Enabled bool `json:"enabled"` - +// ContextCacheSpec configures context caching for a backend. +// +// Setting this field enables context caching: cache_control markers in requests are +// resolved against the shared store at URL, and a cache entry is created when absent. +// Leaving it unset disables caching entirely, and markers are ignored. +type ContextCacheSpec struct { // DefaultTTL is the default time-to-live for newly created cache entries, expressed // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. // Must be at least "60s" per GCP API requirements. @@ -118,4 +116,18 @@ type GCPContextCachingSpec struct { // +optional // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` DefaultTTL string `json:"defaultTTL,omitempty"` + + // URL locates the shared Redis store that records which cachedContents entry a given + // request prefix resolved to. Sharing that state across gateway replicas is what + // prevents two replicas from independently creating a cache for the same prefix. + // + // It is given either as a bare "host:port" + // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. + // + // If Redis is unreachable at request time, caching is skipped and the request is + // served uncached rather than failed. + // + // +kubebuilder:validation:Required + // +kubebuilder:validation:MinLength=1 + URL string `json:"url"` } diff --git a/api/v1beta1/zz_generated.deepcopy.go b/api/v1beta1/zz_generated.deepcopy.go index 3eac1ae533..a979879119 100644 --- a/api/v1beta1/zz_generated.deepcopy.go +++ b/api/v1beta1/zz_generated.deepcopy.go @@ -345,9 +345,9 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.GCPContextCaching != nil { - in, out := &in.GCPContextCaching, &out.GCPContextCaching - *out = new(GCPContextCachingSpec) + if in.ContextCache != nil { + in, out := &in.ContextCache, &out.ContextCache + *out = new(ContextCacheSpec) **out = **in } } @@ -748,6 +748,21 @@ func (in *BackendSecurityPolicyStatus) DeepCopy() *BackendSecurityPolicyStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *ContextCacheSpec) DeepCopyInto(out *ContextCacheSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ContextCacheSpec. +func (in *ContextCacheSpec) DeepCopy() *ContextCacheSpec { + if in == nil { + return nil + } + out := new(ContextCacheSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *CredentialOverrideFromDynamicMetadata) DeepCopyInto(out *CredentialOverrideFromDynamicMetadata) { *out = *in @@ -788,21 +803,6 @@ func (in *CredentialOverrideFromRequestHeaders) DeepCopy() *CredentialOverrideFr return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *GCPContextCachingSpec) DeepCopyInto(out *GCPContextCachingSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new GCPContextCachingSpec. -func (in *GCPContextCachingSpec) DeepCopy() *GCPContextCachingSpec { - if in == nil { - return nil - } - out := new(GCPContextCachingSpec) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *GCPCredentialsFile) DeepCopyInto(out *GCPCredentialsFile) { *out = *in diff --git a/go.mod b/go.mod index 2aa37a2604..c829f33d53 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.42.1 @@ -35,6 +36,7 @@ require ( github.com/prometheus/client_golang v1.24.1 github.com/prometheus/client_model v0.6.2 github.com/prometheus/common v0.70.1 + github.com/redis/go-redis/v9 v9.7.3 github.com/stretchr/testify v1.11.1 github.com/testcontainers/testcontainers-go v0.43.0 github.com/tetratelabs/func-e v1.6.0 @@ -90,6 +92,7 @@ require ( github.com/AzureAD/microsoft-authentication-library-for-go v1.7.2 // indirect github.com/Microsoft/go-winio v0.6.2 // indirect github.com/NYTimes/gziphandler v1.1.1 // indirect + github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 // indirect github.com/antlr4-go/antlr/v4 v4.13.1 // indirect github.com/avast/retry-go/v5 v5.0.0 // indirect github.com/aws/aws-sdk-go-v2/credentials v1.19.28 // indirect @@ -123,6 +126,7 @@ require ( github.com/coreos/go-systemd/v22 v22.7.0 // indirect github.com/cpuguy83/dockercfg v0.3.2 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/distribution/reference v0.6.0 // indirect github.com/docker/cli v29.7.2+incompatible // indirect github.com/docker/docker-credential-helpers v0.9.8 // indirect diff --git a/go.sum b/go.sum index d3fb4f369c..bfd37742c1 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= @@ -385,6 +395,8 @@ github.com/prometheus/otlptranslator v1.0.0 h1:s0LJW/iN9dkIH+EnhiD3BlkkP5QVIUVEo github.com/prometheus/otlptranslator v1.0.0/go.mod h1:vRYWnXvI6aWGpsdY/mOT/cbeVRBlPWtBNDb7kGR3uKM= github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= +github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM= +github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= diff --git a/internal/controller/gateway.go b/internal/controller/gateway.go index 329bcd122c..7ac08dd968 100644 --- a/internal/controller/gateway.go +++ b/internal/controller/gateway.go @@ -524,11 +524,11 @@ func (c *GatewayController) reconcileFilterConfigSecret( b.Schema = schemaToFilterAPI(backendObj.Spec.APISchema) b.HeaderValueFilters = headerValueFiltersToFilterAPI(backendObj.Spec.HeaderValueFilters) - // Wire GCP context caching configuration when present on the backend object. - if backendObj.Spec.GCPContextCaching != nil { - b.GCPContextCaching = &filterapi.GCPContextCaching{ - Enabled: backendObj.Spec.GCPContextCaching.Enabled, - DefaultTTL: backendObj.Spec.GCPContextCaching.DefaultTTL, + // Wire context caching configuration when present on the backend object. + if cc := backendObj.Spec.ContextCache; cc != nil { + b.ContextCache = &filterapi.ContextCache{ + URL: cc.URL, + DefaultTTL: cc.DefaultTTL, } } } diff --git a/internal/controller/gateway_test.go b/internal/controller/gateway_test.go index 974f87d20a..d2c9e9571e 100644 --- a/internal/controller/gateway_test.go +++ b/internal/controller/gateway_test.go @@ -4129,10 +4129,10 @@ func TestGatewayController_warnUndeclaredMetadataNamespaces(t *testing.T) { require.Empty(t, logged) } -// TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching verifies that the -// GCPContextCaching field on an AIServiceBackend is propagated to filterapi.Backend in the +// TestGatewayController_reconcileFilterConfigSecret_ContextCache verifies that the +// ContextCache field on an AIServiceBackend is propagated to filterapi.Backend in the // rendered filter configuration secret. -func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching(t *testing.T) { +func TestGatewayController_reconcileFilterConfigSecret_ContextCache(t *testing.T) { const gwNamespace, someNamespace = "ns", "some-namespace" fakeClient := requireNewFakeClientWithIndexes(t) kube := fake2.NewClientset() @@ -4154,7 +4154,7 @@ func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching(t *test }, } - // Create the AIServiceBackend with GCPContextCaching enabled. + // Create the AIServiceBackend with ContextCache configured. err := fakeClient.Create(t.Context(), &aigv1b1.AIServiceBackend{ ObjectMeta: metav1.ObjectMeta{Name: "gcp-backend", Namespace: gwNamespace}, Spec: aigv1b1.AIServiceBackendSpec{ @@ -4164,9 +4164,9 @@ func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching(t *test Kind: ptr.To(gwapiv1.Kind("Backend")), Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), }, - GCPContextCaching: &aigv1b1.GCPContextCachingSpec{ - Enabled: true, + ContextCache: &aigv1b1.ContextCacheSpec{ DefaultTTL: "600s", + URL: "redis.default.svc.cluster.local:6379", }, }, }) @@ -4186,14 +4186,15 @@ func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching(t *test } } require.NotNil(t, gcpBackend, "expected to find gcp-backend in filter config") - require.NotNil(t, gcpBackend.GCPContextCaching, "GCPContextCaching must be populated in filter config") - require.True(t, gcpBackend.GCPContextCaching.Enabled) - require.Equal(t, "600s", gcpBackend.GCPContextCaching.DefaultTTL) + require.NotNil(t, gcpBackend.ContextCache, "ContextCache must be populated in filter config") + require.Equal(t, "600s", gcpBackend.ContextCache.DefaultTTL) + require.Equal(t, "redis.default.svc.cluster.local:6379", gcpBackend.ContextCache.URL, + "the redis store location must reach the data plane") } -// TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhenAbsent verifies that -// when GCPContextCaching is not set on the AIServiceBackend, filterapi.Backend.GCPContextCaching is nil. -func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhenAbsent(t *testing.T) { +// TestGatewayController_reconcileFilterConfigSecret_ContextCache_NilWhenAbsent verifies that +// when ContextCache is not set on the AIServiceBackend, filterapi.Backend.ContextCache is nil. +func TestGatewayController_reconcileFilterConfigSecret_ContextCache_NilWhenAbsent(t *testing.T) { const gwNamespace, someNamespace = "ns2", "some-namespace2" fakeClient := requireNewFakeClientWithIndexes(t) kube := fake2.NewClientset() @@ -4224,7 +4225,7 @@ func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhen Kind: ptr.To(gwapiv1.Kind("Backend")), Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), }, - // GCPContextCaching intentionally omitted. + // ContextCache intentionally omitted. }, }) require.NoError(t, err) @@ -4242,5 +4243,5 @@ func TestGatewayController_reconcileFilterConfigSecret_GCPContextCaching_NilWhen } } require.NotNil(t, gcpBackend, "expected to find gcp-backend2 in filter config") - require.Nil(t, gcpBackend.GCPContextCaching, "GCPContextCaching must be nil when not set on the CRD") + require.Nil(t, gcpBackend.ContextCache, "ContextCache must be nil when not set on the CRD") } diff --git a/internal/extproc/server.go b/internal/extproc/server.go index be4fcaf5e9..6a954ba184 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -62,6 +62,12 @@ type Server struct { routerProcessorsPerReqID map[string]Processor routerProcessorsPerReqIDMutex sync.RWMutex uuidFn func() string + // cacheResolvers holds one context-cache resolver per GCP backend, keyed by + // "|". Resolvers are carried across config reloads so a + // filterapi update does not tear down and rebuild the Redis connection pool. The + // Redis URL is part of the key so that repointing a backend at a different Redis + // builds a new resolver rather than reusing a pool aimed at the old one. + cacheResolvers map[string]gcpcache.CacheResolver } // NewServer creates a new external processor server. @@ -73,6 +79,7 @@ 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.CacheResolver), uuidFn: uuid.NewString, } return srv, nil @@ -85,16 +92,44 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error return fmt.Errorf("cannot create runtime filter config: %w", err) } - // Attach a shared in-process context-cache resolver for each GCP Vertex AI backend - // that has context caching explicitly enabled. The resolver is reused across requests - // so the in-memory TTL memo is preserved. + // 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 backend that configures + // context caching. Resolvers are reused across reloads so that a config update does + // not tear down and rebuild each backend's Redis connection pool; resolvers for + // backends that disappeared or dropped their cache config are discarded. + live := make(map[string]gcpcache.CacheResolver, len(s.cacheResolvers)) for _, rb := range newConfig.Backends { - if _, ok := rb.Handler.(filterapi.GCPAuthHandler); ok { - if rb.Backend.GCPContextCaching != nil && rb.Backend.GCPContextCaching.Enabled { - rb.CacheResolver = gcpcache.New(nil) + if _, ok := rb.Handler.(filterapi.GCPAuthHandler); !ok { + continue + } + cc := rb.Backend.ContextCache + if cc == nil || cc.URL == "" { + continue + } + key := rb.Backend.Name + "|" + cc.URL + resolver, ok := s.cacheResolvers[key] + if !ok { + store, err := gcpcache.NewRedisStore(cc.URL) + if err != nil { + // A bad URL disables caching for this backend rather than failing the + // reload: the store is an optimization, and rejecting the whole config + // would take down backends that have nothing to do with caching. + logger.Warn("invalid context cache redis url, serving backend uncached", + slog.String("backend", rb.Backend.Name), slog.String("error", err.Error())) + continue } + resolver = gcpcache.New(nil, store, logger) } + live[key] = resolver + rb.CacheResolver = resolver } + s.cacheResolvers = live s.config = newConfig // This is racey, but we don't care. return nil diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index da6cf04314..a125759721 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -54,7 +54,7 @@ func TestServer_LoadConfig(t *testing.T) { } // TestServer_LoadConfig_GCPCacheResolver verifies that a CacheResolver is attached -// to a GCP backend only when GCPContextCaching.Enabled is true. +// to a GCP backend only when ContextCache is configured with a store URL. func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { s := &Server{} @@ -66,15 +66,15 @@ func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { Auth: &filterapi.BackendAuth{ GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, }, - GCPContextCaching: &filterapi.GCPContextCaching{Enabled: true, DefaultTTL: "600s"}, + ContextCache: &filterapi.ContextCache{URL: "localhost:6379", DefaultTTL: "600s"}, }, { - Name: "gcp-caching-off", + Name: "gcp-caching-no-url", Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, Auth: &filterapi.BackendAuth{ GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, }, - GCPContextCaching: &filterapi.GCPContextCaching{Enabled: false}, + ContextCache: &filterapi.ContextCache{}, }, { Name: "gcp-no-caching-field", @@ -82,7 +82,7 @@ func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { Auth: &filterapi.BackendAuth{ GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, }, - // GCPContextCaching is nil. + // ContextCache is nil. }, { Name: "non-gcp", @@ -96,18 +96,90 @@ func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { require.NoError(t, err) rb := s.config.Backends["gcp-caching-on"] - require.NotNil(t, rb.CacheResolver, "CacheResolver must be set when GCPContextCaching.Enabled=true") + require.NotNil(t, rb.CacheResolver, "CacheResolver must be set when ContextCache has a store URL") - rb = s.config.Backends["gcp-caching-off"] - require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when GCPContextCaching.Enabled=false") + rb = s.config.Backends["gcp-caching-no-url"] + require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when ContextCache has no store URL") rb = s.config.Backends["gcp-no-caching-field"] - require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when GCPContextCaching is absent") + require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when ContextCache is absent") rb = s.config.Backends["non-gcp"] require.Nil(t, rb.CacheResolver, "CacheResolver must be nil for non-GCP backends") } +// gcpBackendWithCaching builds a GCP backend whose context caching points at redisURL. +func gcpBackendWithCaching(name, redisURL string) filterapi.Backend { + return filterapi.Backend{ + Name: name, + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}}, + ContextCache: &filterapi.ContextCache{URL: redisURL, DefaultTTL: "600s"}, + } +} + +// A resolver owns a Redis connection pool. Rebuilding it on every filterapi update would +// discard live connections on a hot path that fires whenever any unrelated backend changes. +func TestServer_LoadConfig_GCPCacheResolver_ReusedAcrossReloads(t *testing.T) { + s := &Server{} + config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}} + + require.NoError(t, s.LoadConfig(t.Context(), config)) + first := s.config.Backends["gcp"].CacheResolver + require.NotNil(t, first) + + require.NoError(t, s.LoadConfig(t.Context(), config)) + require.Same(t, first, s.config.Backends["gcp"].CacheResolver, + "an unchanged backend must keep its resolver, and with it its connection pool") +} + +// Repointing a backend at a different Redis must not keep reusing a pool dialing the old one. +func TestServer_LoadConfig_GCPCacheResolver_RebuiltWhenRedisURLChanges(t *testing.T) { + s := &Server{} + + require.NoError(t, s.LoadConfig(t.Context(), + &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}})) + first := s.config.Backends["gcp"].CacheResolver + require.NotNil(t, first) + + require.NoError(t, s.LoadConfig(t.Context(), + &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6380")}})) + require.NotSame(t, first, s.config.Backends["gcp"].CacheResolver, + "a new Redis URL must produce a new resolver") +} + +// Resolvers for backends that vanish must not accumulate: the map is rebuilt each reload. +func TestServer_LoadConfig_GCPCacheResolver_DroppedWhenBackendDisappears(t *testing.T) { + s := &Server{} + + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp-a", "localhost:6379"), + gcpBackendWithCaching("gcp-b", "localhost:6379"), + }})) + require.Len(t, s.cacheResolvers, 2) + + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp-a", "localhost:6379"), + }})) + require.Len(t, s.cacheResolvers, 1, "the departed backend's resolver must be released") +} + +// A malformed Redis URL disables caching for that backend rather than failing the whole +// reload, which would take down backends that have nothing to do with caching. +func TestServer_LoadConfig_GCPCacheResolver_BadRedisURLDoesNotFailReload(t *testing.T) { + s := &Server{} + config := &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp-bad-redis", "http://not-a-redis-url"), + gcpBackendWithCaching("gcp-ok", "localhost:6379"), + }} + + require.NoError(t, s.LoadConfig(t.Context(), config), "a bad redis url must not fail the reload") + require.Nil(t, s.config.Backends["gcp-bad-redis"].CacheResolver, + "a backend whose store cannot be built is served uncached") + require.NotNil(t, s.config.Backends["gcp-ok"].CacheResolver, + "an unrelated backend must keep its resolver") +} + func TestServer_Check(t *testing.T) { s, _ := requireNewServerWithMockProcessor(t) diff --git a/internal/filterapi/filterconfig.go b/internal/filterapi/filterconfig.go index f6853c2712..1aec42eda9 100644 --- a/internal/filterapi/filterconfig.go +++ b/internal/filterapi/filterconfig.go @@ -207,22 +207,27 @@ type Backend struct { // HeaderValueFilters filter individual values out of multi-valued request headers before sending // the request to the backend. Optional. HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` - // GCPContextCaching configures Gemini context caching for this backend. Optional. - // When non-nil and Enabled is true, the extproc will resolve or create a Google - // cachedContents entry for requests that carry cache_control markers. - GCPContextCaching *GCPContextCaching `json:"gcpContextCaching,omitempty"` + // ContextCache configures context caching for this backend. Optional. + // When non-nil, the extproc will resolve or create a Google cachedContents entry + // for requests that carry cache_control markers. + ContextCache *ContextCache `json:"contextCache,omitempty"` } -// GCPContextCaching configures in-process Gemini context caching for a GCP Vertex AI backend. -type GCPContextCaching struct { - // Enabled controls whether context caching is active for this backend. - // When false (or when this struct is nil), cache_control markers in requests - // are ignored and no cachedContents API calls are made. - Enabled bool `json:"enabled"` +// ContextCache configures context caching for a GCP Vertex AI backend. +// +// Its presence enables caching: markers are resolved against the store at URL and a +// cache entry is created when absent. +type ContextCache struct { // DefaultTTL is the default time-to-live for newly created cache entries expressed // as a GCP duration string (e.g. "600s"). When empty, the resolver default (300s) is used. // +optional DefaultTTL string `json:"defaultTTL,omitempty"` + // URL points at the shared Redis store that records which cachedContents entry a + // given prefix resolved to. It is what stops two gateway replicas from + // independently creating a cache for the same prefix. + // + // It is either a bare "host:port" or a full "redis://" URL. + URL string `json:"url"` } // BackendAuth corresponds partially to BackendSecurityPolicy in api/v1alpha1/api.go. diff --git a/internal/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go index 20b7f90c5b..695b55aa52 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/gcpcache/gcpcache.go @@ -3,18 +3,21 @@ // The full text of the Apache license is available in the LICENSE file at // the root of the repo. -// Package gcpcache implements in-process context-cache resolution for GCP Vertex AI (Gemini). +// 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 up an in-memory TTL memo. -// 4. On a memo miss, lists Google cachedContents for the target region and model. +// 3. Generates a deterministic SHA-256 cache key and looks it up in the shared store. +// 4. On a store miss, lists Google cachedContents for the target region and model. // 5. Creates the cache entry if not found, then returns the cache resource name and the // non-cached messages. // +// Concurrent resolutions of the same key are deduplicated by the store's create lock, +// which coordinates across gateway replicas rather than only within one process. +// // The CacheResolver interface allows the implementation to be replaced with an external // cache service in the future without changing the request-path callers. package gcpcache @@ -26,8 +29,8 @@ import ( "encoding/hex" "fmt" "io" + "log/slog" "net/http" - "sync" "time" "google.golang.org/genai" @@ -45,6 +48,11 @@ const ( // gcpCachedContentsBasePath is the base URL for the Vertex AI cachedContents REST API. gcpCachedContentsBasePath = "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/cachedContents" + + // staleThreshold is how close to its expiry a stored entry must be before it is + // treated as a miss, so there is time to re-resolve before the Google cache + // disappears underneath an in-flight request. + staleThreshold = 10 * time.Second ) // ResolveResult holds the outcome of a successful cache resolution. @@ -74,37 +82,40 @@ type CacheResolver interface { Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) } -// memoEntry is an in-memory cache entry for resolved Google cache names. -type memoEntry struct { - cacheName string - expireTime time.Time -} - -// resolver is the default in-process CacheResolver implementation. +// resolver is the default CacheResolver implementation. type resolver struct { httpClient *http.Client - mu sync.Mutex - memo map[string]memoEntry // key → memoEntry; guarded by mu + // store holds resolved cache names. It is shared across replicas when backed by + // Redis, which is what keeps two replicas from creating the same cache. Its + // failures are never fatal: they are logged and treated as a miss. + store CacheStore + + // logger records store failures, which are otherwise invisible because the request + // proceeds normally. + logger *slog.Logger } -// New creates a new CacheResolver. The provided httpClient is used for calls to the -// Google cachedContents API; pass nil to use the default http.Client. -func New(httpClient *http.Client) CacheResolver { +// New creates a new CacheResolver. +// +// httpClient is used for calls to the Google 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. logger may be nil. +func New(httpClient *http.Client, store CacheStore, logger *slog.Logger) CacheResolver { if httpClient == nil { httpClient = &http.Client{Timeout: 30 * time.Second} } - return &resolver{ - httpClient: httpClient, - memo: make(map[string]memoEntry), + if store == nil { + store = noopStore{} + } + if logger == nil { + logger = slog.New(slog.DiscardHandler) } + return &resolver{httpClient: httpClient, store: store, logger: logger} } // Resolve implements CacheResolver. func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) { - region := gcpAuth.GCPRegion() - project := gcpAuth.GCPProject() - // Find the last cache_control breakpoint in the message list. breakpoint := findBreakpoint(openAIReq.Messages) if breakpoint < 0 { @@ -135,15 +146,51 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion return nil, fmt.Errorf("gcpcache: failed to compute cache key: %w", err) } - // Check in-memory memo first. - if entry, ok := r.getMemo(cacheKey); ok { + // 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. Google-side failures below are a different matter and do fail fast. + if e, ok := r.storeGet(ctx, cacheKey); ok { return &ResolveResult{ - CacheName: entry.cacheName, + CacheName: e.cacheName, Messages: remainderMessages, - ExpireTime: entry.expireTime, + ExpireTime: e.expireTime, }, nil } + // Store miss — resolve against GCP. Concurrent resolutions of the same key are + // arbitrated by the store's create lock (see resolveUncached), which deduplicates + // across replicas rather than only within this one. + result, err := r.resolveUncached(ctx, openAIReq, gcpAuth, cacheKey, ttl, contents, systemInstruction, geminiTools) + if err != nil { + return nil, err + } + + // resolveUncached leaves Messages unset because it deals only in the cached prefix; + // the remainder is per-request. + result.Messages = remainderMessages + return result, nil +} + +// resolveUncached performs the GCP-side resolution for a cache key: list, create if +// absent, then re-list to converge duplicate-create races. It deals only in the cached +// prefix, so it leaves ResolveResult.Messages unset for the caller to fill in per-request. +// +// Before creating, it claims a store-level lock so that replicas racing on a cold prefix +// produce one create rather than N. Losing the race means waiting for the winner's result; +// the wait degrading into a plain resolution is safe, because the list below is the actual +// source of truth. +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, +) (*ResolveResult, error) { + region := gcpAuth.GCPRegion() + project := gcpAuth.GCPProject() + // Get access token. tokenSrc := gcpAuth.GCPTokenSource() token, err := tokenSrc.Token() @@ -152,25 +199,43 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion } accessToken := token.AccessToken + // Claim the right to create this key. A replica that loses waits for the winner's + // result; if the winner never publishes one, it falls through and resolves itself. + locked := r.tryLock(ctx, cacheKey) + if !locked { + if e, ok := r.awaitLeader(ctx, cacheKey); ok { + return &ResolveResult{CacheName: e.cacheName, ExpireTime: e.expireTime}, nil + } + } + // List existing caches and match by displayName (cacheKey). baseURL := fmt.Sprintf(gcpCachedContentsBasePath, region, project, region) existingName, expireTime, err := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) if err != nil { + if locked { + r.unlock(ctx, cacheKey) + } return nil, fmt.Errorf("gcpcache: failed to list cached contents: %w", err) } if existingName != "" { - r.setMemo(cacheKey, existingName, expireTime) + r.storeSet(ctx, cacheKey, entry{cacheName: existingName, expireTime: expireTime}) return &ResolveResult{ CacheName: existingName, - Messages: remainderMessages, ExpireTime: expireTime, }, nil } - // Cache not found — create it. + // Cache not found — create it. A replica that did not win the lock still creates + // here, having already waited for the winner without result; this is the duplicate + // the lock narrows but cannot fully eliminate. created, tokenCount, expireTime, err := r.createCache(ctx, baseURL, accessToken, openAIReq.Model, region, project, cacheKey, contents, systemInstruction, geminiTools, ttl) if err != nil { + // Release the claim so waiters retry immediately rather than blocking for the + // remainder of the lock TTL on a create that will never publish. + if locked { + r.unlock(ctx, cacheKey) + } return nil, fmt.Errorf("gcpcache: failed to create cached content: %w", err) } @@ -178,24 +243,112 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion finalName, finalExpire, listErr := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) if listErr == nil && finalName != "" && finalName != created { // Another replica created a cache with the same key; use the one from the list. - r.setMemo(cacheKey, finalName, finalExpire) + r.storeSet(ctx, cacheKey, entry{cacheName: finalName, expireTime: finalExpire}) return &ResolveResult{ CacheName: finalName, - Messages: remainderMessages, ExpireTime: finalExpire, }, nil } - r.setMemo(cacheKey, created, expireTime) + r.storeSet(ctx, cacheKey, entry{cacheName: created, expireTime: expireTime}) return &ResolveResult{ CacheName: created, - Messages: remainderMessages, 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 Google as the source of +// truth. Google-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) (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 entry{}, false + } + if !ok { + return 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) < staleThreshold { + return entry{}, false + } + return e, true +} + +// storeSet records a resolved cache name, expiring it with the Google entry itself so +// the store cannot outlive what it points at. +func (r *resolver) storeSet(ctx context.Context, key string, e 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())) + } +} + +// locker is implemented by stores that can arbitrate cache creation across replicas. +// Stores that cannot (such as the no-op store) simply never grant a lock, which leaves +// every caller resolving independently — correct, just without the deduplication. +type locker interface { + tryLock(ctx context.Context, key string) (bool, error) + unlock(ctx context.Context, key string) + awaitLeader(ctx context.Context, key string) (entry, bool, error) +} + +// tryLock claims the right to create key. It reports false when the store cannot lock, +// when the lock is held elsewhere, or on failure — all cases where the caller should +// resolve against Google rather than wait. +func (r *resolver) tryLock(ctx context.Context, key string) bool { + l, ok := r.store.(locker) + if !ok { + return false + } + won, err := l.tryLock(ctx, key) + if err != nil { + r.logger.Warn("gcpcache: cache store lock failed, proceeding without coordination", + slog.String("error", err.Error())) + return false + } + return won +} + +func (r *resolver) unlock(ctx context.Context, key string) { + if l, ok := r.store.(locker); ok { + l.unlock(ctx, key) + } +} + +// awaitLeader waits for the replica holding the lock to publish its result. +func (r *resolver) awaitLeader(ctx context.Context, key string) (entry, bool) { + l, ok := r.store.(locker) + if !ok { + return entry{}, false + } + e, found, err := l.awaitLeader(ctx, key) + if err != nil { + // Includes the timeout case: fall through and resolve rather than fail. + return entry{}, false + } + return e, found +} + // ----------------------------------------------------------------------- // Breakpoint and TTL helpers // ----------------------------------------------------------------------- @@ -347,31 +500,6 @@ func computeCacheKey(model string, contents []genai.Content, systemInstruction * return hex.EncodeToString(h[:]), nil } -// ----------------------------------------------------------------------- -// In-memory memo -// ----------------------------------------------------------------------- - -func (r *resolver) getMemo(key string) (memoEntry, bool) { - r.mu.Lock() - defer r.mu.Unlock() - entry, ok := r.memo[key] - if !ok { - return memoEntry{}, false - } - // Treat entries expiring within 10 s as stale so we have time to re-resolve. - if time.Until(entry.expireTime) < 10*time.Second { - delete(r.memo, key) - return memoEntry{}, false - } - return entry, true -} - -func (r *resolver) setMemo(key, cacheName string, expireTime time.Time) { - r.mu.Lock() - defer r.mu.Unlock() - r.memo[key] = memoEntry{cacheName: cacheName, expireTime: expireTime} -} - // ----------------------------------------------------------------------- // Google cachedContents REST API calls // ----------------------------------------------------------------------- diff --git a/internal/gcpcache/gcpcache_test.go b/internal/gcpcache/gcpcache_test.go index 828d719809..e0a0141f1f 100644 --- a/internal/gcpcache/gcpcache_test.go +++ b/internal/gcpcache/gcpcache_test.go @@ -9,6 +9,8 @@ import ( "context" "net/http" "net/http/httptest" + "sync" + "sync/atomic" "testing" "time" @@ -97,10 +99,15 @@ type fakeCacheServer struct { listResponse string createStatus int createBody string - listCalls int - createCalls int + // 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{ @@ -112,11 +119,11 @@ func newFakeCacheServer(t *testing.T, listBody, createBody string, createStatus w.Header().Set("Content-Type", "application/json") switch r.Method { case http.MethodGet: - f.listCalls++ + f.listCalls.Add(1) w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(f.listResponse)) case http.MethodPost: - f.createCalls++ + f.createCalls.Add(1) w.WriteHeader(f.createStatus) _, _ = w.Write([]byte(f.createBody)) } @@ -137,12 +144,62 @@ func (t *redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) return http.DefaultTransport.RoundTrip(req2) } -func resolverWithServer(srvURL string) *resolver { +// 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 ...CacheStore) *resolver { host := srvURL[len("http://"):] + var s CacheStore + if len(store) > 0 { + s = store[0] + } return New(&http.Client{ Transport: &redirectTransport{fakeHost: host}, Timeout: 5 * time.Second, - }).(*resolver) + }, s, nil).(*resolver) +} + +// memStore is an in-process CacheStore for tests. It implements only the CacheStore +// contract, not locker, so resolvers using it never coordinate — which is what the +// no-op-store path does in production too. +type memStore struct { + mu sync.Mutex + entries map[string]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]entry{}} } + +func (m *memStore) Get(_ context.Context, key string) (entry, bool, error) { + m.gets.Add(1) + if m.getErr != nil { + return 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 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) seed(key, cacheName string, expireTime time.Time) { + m.mu.Lock() + defer m.mu.Unlock() + m.entries[key] = entry{cacheName: cacheName, expireTime: expireTime} } // computeKeyForRequest replicates the key generation for use in test assertions. @@ -293,7 +350,7 @@ func TestComputeCacheKey_DifferentModels_DifferentKeys(t *testing.T) { // ----------------------------------------------------------------------- func TestResolver_NoMarkers_ReturnsNil(t *testing.T) { - r := New(nil).(*resolver) + r := New(nil, nil, nil).(*resolver) req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{userMsg("hello")}, @@ -328,8 +385,8 @@ func TestResolver_CacheMiss_Creates(t *testing.T) { // remainder = messages after the breakpoint (index 0 → only userMsg at index 1 remains) assert.Equal(t, []openai.ChatCompletionMessageParamUnion{userMsg("Hello")}, res.Messages) // list → create → re-list (post-create convergence check) - assert.Equal(t, 2, fake.listCalls) - assert.Equal(t, 1, fake.createCalls) + assert.Equal(t, 2, fake.lists()) + assert.Equal(t, 1, fake.creates()) } func TestResolver_CacheHit_FromGoogleList(t *testing.T) { @@ -364,12 +421,13 @@ func TestResolver_CacheHit_FromGoogleList(t *testing.T) { require.NotNil(t, res) assert.Equal(t, "projects/p/locations/us-central1/cachedContents/existing", res.CacheName) assert.False(t, res.Created) - assert.Equal(t, 0, fake.createCalls) + assert.Equal(t, 0, fake.creates()) } -func TestResolver_MemoHit_SkipsGoogleAPICalls(t *testing.T) { +func TestResolver_StoreHit_SkipsGoogleAPICalls(t *testing.T) { fake := newFakeCacheServer(t, `{"cachedContents":[]}`, "", http.StatusOK) - r := resolverWithServer(fake.srv.URL) + store := newMemStore() + r := resolverWithServer(fake.srv.URL, store) req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", @@ -379,19 +437,19 @@ func TestResolver_MemoHit_SkipsGoogleAPICalls(t *testing.T) { }, } key := computeKeyForRequest(t, req) - r.setMemo(key, "projects/p/locations/r/cachedContents/memo-hit", time.Now().Add(10*time.Minute)) + store.seed(key, "projects/p/locations/r/cachedContents/store-hit", time.Now().Add(10*time.Minute)) auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} res, err := r.Resolve(context.Background(), req, auth) require.NoError(t, err) require.NotNil(t, res) - assert.Equal(t, "projects/p/locations/r/cachedContents/memo-hit", res.CacheName) + assert.Equal(t, "projects/p/locations/r/cachedContents/store-hit", res.CacheName) assert.False(t, res.Created) - assert.Equal(t, 0, fake.listCalls, "memo hit must not call the Google API") - assert.Equal(t, 0, fake.createCalls) + assert.Equal(t, 0, fake.lists(), "store hit must not call the Google API") + assert.Equal(t, 0, fake.creates()) } -func TestResolver_MemoExpiry_Refetches(t *testing.T) { +func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -413,16 +471,17 @@ func TestResolver_MemoExpiry_Refetches(t *testing.T) { }, }) fake := newFakeCacheServer(t, string(listBody), "", http.StatusOK) - r := resolverWithServer(fake.srv.URL) - // Seed with entry expiring within the 10 s stale window → should be evicted. - r.setMemo(key, "projects/p/locations/us-central1/cachedContents/stale", time.Now().Add(5*time.Second)) + 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)) auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} res, err := r.Resolve(context.Background(), req, auth) 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.listCalls, "stale memo must trigger a list call") + assert.Equal(t, 1, fake.lists(), "a near-expiry store entry must trigger a list call") } func TestResolver_CreateFailure_ReturnsError(t *testing.T) { diff --git a/internal/gcpcache/redis.go b/internal/gcpcache/redis.go new file mode 100644 index 0000000000..12b521f50e --- /dev/null +++ b/internal/gcpcache/redis.go @@ -0,0 +1,190 @@ +// 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" + "errors" + "fmt" + "strings" + "time" + + "github.com/redis/go-redis/v9" +) + +const ( + // lockSentinel marks a key as claimed by a replica that is creating the cache. + // A reader seeing it knows a create is in flight rather than that a name is stored. + lockSentinel = "\x00creating" + + // lockTTL bounds how long a create lock is held. It must exceed the p99 duration of + // a Google cachedContents create, or a second replica will claim the lock while the + // first is still working and both will create. It is deliberately not configurable: + // it is internal to the lock protocol, and widening the config surface later is + // additive whereas removing a field is a breaking change. + lockTTL = 30 * time.Second + + // lockWaitTime bounds how long a replica polls for the leader's result before giving + // up and resolving against Google itself. It must stay well under Envoy's ext_proc + // message timeout: a waiting request is blocked inside ProcessRequestHeaders. + lockWaitTime = 5 * time.Second + + // lockPollInterval is how often a waiter re-reads the key while the leader works. + lockPollInterval = 50 * time.Millisecond +) + +// errLockHeld reports that another replica holds the create lock and did not publish a +// result before lockWaitTime elapsed. +var errLockHeld = errors.New("gcpcache: create lock held by another replica") + +// redisStore is a CacheStore backed by Redis, shared across gateway replicas. +// +// Beyond plain get/set it arbitrates cache creation: the first replica to claim a key +// creates the Google cache while the others wait for its result, so a cold prefix +// arriving at N replicas produces one create rather than N. +type redisStore struct { + client redis.UniversalClient +} + +// NewRedisStore returns a CacheStore 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 the resolver treats as a miss. +func NewRedisStore(url string) (CacheStore, error) { + opts, err := parseRedisURL(url) + if err != nil { + return nil, err + } + return &redisStore{client: redis.NewClient(opts)}, nil +} + +// parseRedisURL accepts both a scheme-qualified URL and a bare host:port. +func parseRedisURL(url string) (*redis.Options, error) { + if url == "" { + return nil, errors.New("gcpcache: redis url is empty") + } + if strings.Contains(url, "://") { + opts, err := redis.ParseURL(url) + if err != nil { + return nil, fmt.Errorf("gcpcache: invalid redis url: %w", err) + } + return opts, nil + } + return &redis.Options{Addr: url}, nil +} + +// Get implements CacheStore. A key holding the create sentinel is reported as a miss, +// since no cache name is available yet. +func (s *redisStore) Get(ctx context.Context, key string) (entry, bool, error) { + v, err := s.client.Get(ctx, key).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + return entry{}, false, nil + } + return entry{}, false, fmt.Errorf("gcpcache: redis get: %w", err) + } + if v == lockSentinel { + return entry{}, false, nil + } + 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 entry{}, false, nil + } + return e, true, nil +} + +// Set implements CacheStore. +func (s *redisStore) Set(ctx context.Context, key string, e entry, ttl time.Duration) error { + if err := s.client.Set(ctx, key, encodeEntry(e), ttl).Err(); err != nil { + return fmt.Errorf("gcpcache: redis set: %w", err) + } + return nil +} + +// tryLock attempts to claim key for creation. It reports whether the caller won. +// +// The winner must publish its result with Set (or release the claim with unlock) so that +// waiters are not stuck until the lock expires. +func (s *redisStore) tryLock(ctx context.Context, key string) (bool, error) { + won, err := s.client.SetNX(ctx, key, lockSentinel, lockTTL).Result() + if err != nil { + return false, fmt.Errorf("gcpcache: redis lock: %w", err) + } + return won, nil +} + +// unlock releases a claim that did not produce a cache name, so waiters retry promptly +// instead of blocking for the remainder of lockTTL. +// +// It deletes only a key still holding the sentinel: if a create succeeded and published a +// real name, or the lock expired and another replica re-claimed it, the key is left alone. +func (s *redisStore) unlock(ctx context.Context, key string) { + const script = `if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) end return 0` + // Best-effort: a failure here only means waiters fall back to Google after lockWaitTime. + _ = s.client.Eval(ctx, script, []string{key}, lockSentinel).Err() +} + +// awaitLeader polls key until the leader publishes a cache name, the claim disappears, or +// lockWaitTime elapses. +// +// A timeout returns errLockHeld, which the caller treats as "resolve against Google +// yourself". That degrades a stuck leader into an extra list rather than a duplicate +// create, since the resolver lists before it creates. +func (s *redisStore) awaitLeader(ctx context.Context, key string) (entry, bool, error) { + deadline := time.NewTimer(lockWaitTime) + defer deadline.Stop() + ticker := time.NewTicker(lockPollInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return entry{}, false, ctx.Err() + case <-deadline.C: + return entry{}, false, errLockHeld + case <-ticker.C: + v, err := s.client.Get(ctx, key).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + // The leader released without publishing; caller should resolve. + return entry{}, false, errLockHeld + } + return entry{}, false, fmt.Errorf("gcpcache: redis get: %w", err) + } + if v == lockSentinel { + continue // Still creating. + } + e, decErr := decodeEntry(v) + if decErr != nil { + return entry{}, false, errLockHeld + } + return e, true, 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 entry) string { + return e.expireTime.UTC().Format(time.RFC3339) + "|" + e.cacheName +} + +func decodeEntry(s string) (entry, error) { + expiry, name, ok := strings.Cut(s, "|") + if !ok || name == "" { + return entry{}, fmt.Errorf("gcpcache: malformed cache entry %q", s) + } + t, err := time.Parse(time.RFC3339, expiry) + if err != nil { + return entry{}, fmt.Errorf("gcpcache: malformed cache entry expiry %q: %w", expiry, err) + } + return entry{cacheName: name, expireTime: t}, nil +} + +var _ CacheStore = (*redisStore)(nil) diff --git a/internal/gcpcache/redis_test.go b/internal/gcpcache/redis_test.go new file mode 100644 index 0000000000..f81ef3b135 --- /dev/null +++ b/internal/gcpcache/redis_test.go @@ -0,0 +1,312 @@ +// 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" + "net/http" + "sync" + "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" +) + +// 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 := NewRedisStore(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", entry{cacheName: "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.cacheName) + 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 key holding the create sentinel is a claim, not a result: reporting it as a hit would +// hand the caller the sentinel string as a cache name. +func TestRedisStore_SentinelIsAMiss(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + + _, ok, err := s.Get(ctx, "k") + require.NoError(t, err) + assert.False(t, ok, "an in-flight create must not read as a cache hit") +} + +// 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 TestRedisStore_TryLockIsExclusive(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + + first, err := s.tryLock(ctx, "k") + require.NoError(t, err) + assert.True(t, first) + + second, err := s.tryLock(ctx, "k") + require.NoError(t, err) + assert.False(t, second, "a second claim on a held key must lose") +} + +// unlock must not delete a key that has moved on from the sentinel, or it would erase a +// published cache name (or another replica's fresh claim after the lock expired). +func TestRedisStore_UnlockOnlyRemovesOwnSentinel(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + expire := time.Now().Add(10 * time.Minute) + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + require.NoError(t, s.Set(ctx, "k", entry{cacheName: "published", expireTime: expire}, time.Minute)) + + s.unlock(ctx, "k") + + got, ok, err := s.Get(ctx, "k") + require.NoError(t, err) + require.True(t, ok, "unlock must not erase a published result") + assert.Equal(t, "published", got.cacheName) +} + +func TestRedisStore_UnlockReleasesClaim(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + s.unlock(ctx, "k") + + again, err := s.tryLock(ctx, "k") + require.NoError(t, err) + assert.True(t, again, "a released claim must be re-claimable immediately") +} + +func TestRedisStore_AwaitLeader_ReturnsPublishedResult(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + expire := time.Now().Add(10 * time.Minute).UTC().Truncate(time.Second) + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + + go func() { + time.Sleep(100 * time.Millisecond) + _ = s.Set(ctx, "k", entry{cacheName: "from-leader", expireTime: expire}, time.Minute) + }() + + got, ok, err := s.awaitLeader(ctx, "k") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "from-leader", got.cacheName) +} + +// A leader that releases without publishing must not strand its waiters for the rest of +// lockWaitTime; they are told to resolve themselves as soon as the claim disappears. +func TestRedisStore_AwaitLeader_ClaimReleasedWithoutResult(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx := context.Background() + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + + go func() { + time.Sleep(100 * time.Millisecond) + s.unlock(ctx, "k") + }() + + start := time.Now() + _, ok, err := s.awaitLeader(ctx, "k") + require.ErrorIs(t, err, errLockHeld) + assert.False(t, ok) + assert.Less(t, time.Since(start), lockWaitTime, "waiter must not block for the full wait window") +} + +func TestRedisStore_AwaitLeader_HonorsContextCancellation(t *testing.T) { + s, _ := newTestRedisStore(t) + ctx, cancel := context.WithCancel(context.Background()) + + won, err := s.tryLock(ctx, "k") + require.NoError(t, err) + require.True(t, won) + + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + + _, _, err = s.awaitLeader(ctx, "k") + require.ErrorIs(t, err, context.Canceled) +} + +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 := entry{cacheName: "projects/p|weird/cachedContents/x", expireTime: expire} + + got, err := decodeEntry(encodeEntry(e)) + require.NoError(t, err) + assert.Equal(t, e.cacheName, got.cacheName) + 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) + } +} + +// ----------------------------------------------------------------------- +// Cross-replica behavior +// ----------------------------------------------------------------------- + +// Two resolvers stand in for two gateway replicas: separate processes, one shared +// Redis. Deduplicating a cold prefix across them is the reason the store's create +// lock exists. +func TestResolver_CrossReplica_SingleCreate(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 := NewRedisStore(mr.Addr()) + require.NoError(t, err) + return resolverWithServer(fake.srv.URL, store) + } + a, b := newReplica(), newReplica() + + req := func() *openai.ChatCompletionRequest { + return &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, + } + } + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + + var start, done sync.WaitGroup + start.Add(1) + results := make([]*ResolveResult, 2) + errs := make([]error, 2) + for i, r := range []*resolver{a, b} { + done.Add(1) + go func() { + defer done.Done() + start.Wait() + results[i], errs[i] = r.Resolve(context.Background(), req(), auth) + }() + } + start.Done() + done.Wait() + + for i := range results { + require.NoError(t, errs[i], "replica %d", i) + require.NotNil(t, results[i], "replica %d", i) + assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new", results[i].CacheName) + } + assert.Equal(t, 1, fake.creates(), "two replicas sharing one Redis must create exactly one cache") + + // Only the replica that performed the write may bill for it. + created := 0 + for _, res := range results { + if res.Created { + created++ + } + } + assert.Equal(t, 1, created, "exactly one replica may report Created") +} + +// 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 := NewRedisStore(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"), + }, + } + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + + res, err := r.Resolve(context.Background(), req, auth) + 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/gcpcache/store.go b/internal/gcpcache/store.go new file mode 100644 index 0000000000..e643389b2f --- /dev/null +++ b/internal/gcpcache/store.go @@ -0,0 +1,48 @@ +// 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" + "time" +) + +// entry is a resolved Google cachedContents entry: the resource name and when it expires. +type entry struct { + cacheName string + expireTime time.Time +} + +// CacheStore is the storage layer behind the resolver, mapping a deterministic cache +// key to the Google cache name it resolved to. +// +// The store is shared across gateway replicas, which is what prevents two replicas +// from independently creating a cache for the same prefix. Implementations must be +// safe for concurrent use. +// +// Store errors are never fatal to a request: the resolver treats them as a cache miss +// and proceeds to Google. See the failure policy in the Redis cache store proposal. +type CacheStore 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 +} + +// noopStore is the CacheStore 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. +type noopStore struct{} + +func (noopStore) Get(context.Context, string) (entry, bool, error) { return entry{}, false, nil } + +func (noopStore) Set(context.Context, string, entry, time.Duration) error { return nil } + +var _ CacheStore = noopStore{} diff --git a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml index 52ac704620..ef7342939b 100644 --- a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml +++ b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml @@ -224,13 +224,13 @@ spec: - path x-kubernetes-list-type: map type: object - gcpContextCaching: + contextCache: description: |- - GCPContextCaching configures Gemini context caching for this backend. - When set, the ai-gateway will automatically resolve or create a Google cachedContents + ContextCache configures context caching for this backend. + When set, the ai-gateway will automatically resolve or create a context cache entry for requests that carry Anthropic-style cache_control markers, and inject the - resolved cache resource name into the Gemini API request. - This field is only meaningful for GCPVertexAI backends. + resolved cache resource name into the API request. + Currently only supported for GCP Vertex AI backends. properties: defaultTTL: description: |- @@ -239,14 +239,21 @@ spec: Must be at least "60s" per GCP API requirements. pattern: ^[1-9][0-9]*s$ type: string - enabled: + url: description: |- - Enabled controls whether context caching is active for this backend. - When false, cache_control markers in requests are ignored and no cachedContents - API calls are made. - type: boolean + URL locates the shared Redis store that records which cachedContents entry a given + request prefix resolved to. Sharing that state across gateway replicas is what + prevents two replicas from independently creating a cache for the same prefix. + + It is given either as a bare "host:port" + (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. + + If Redis is unreachable at request time, caching is skipped and the request is + served uncached rather than failed. + minLength: 1 + type: string required: - - enabled + - url type: object headerMutation: description: |- @@ -665,13 +672,13 @@ spec: - path x-kubernetes-list-type: map type: object - gcpContextCaching: + contextCache: description: |- - GCPContextCaching configures Gemini context caching for this backend. - When set, the ai-gateway will automatically resolve or create a Google cachedContents + ContextCache configures context caching for this backend. + When set, the ai-gateway will automatically resolve or create a context cache entry for requests that carry Anthropic-style cache_control markers, and inject the - resolved cache resource name into the Gemini API request. - This field is only meaningful for GCPVertexAI backends. + resolved cache resource name into the API request. + Currently only supported for GCP Vertex AI backends. properties: defaultTTL: description: |- @@ -680,14 +687,21 @@ spec: Must be at least "60s" per GCP API requirements. pattern: ^[1-9][0-9]*s$ type: string - enabled: + url: description: |- - Enabled controls whether context caching is active for this backend. - When false, cache_control markers in requests are ignored and no cachedContents - API calls are made. - type: boolean + URL locates the shared Redis store that records which cachedContents entry a given + request prefix resolved to. Sharing that state across gateway replicas is what + prevents two replicas from independently creating a cache for the same prefix. + + It is given either as a bare "host:port" + (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. + + If Redis is unreachable at request time, caching is skipped and the request is + served uncached rather than failed. + minLength: 1 + type: string required: - - enabled + - url type: object headerMutation: description: |- diff --git a/site/docs/api/api.mdx b/site/docs/api/api.mdx index 47791b5711..b3d3cde32a 100644 --- a/site/docs/api/api.mdx +++ b/site/docs/api/api.mdx @@ -579,7 +579,7 @@ QuotaPolicyList contains a list of QuotaPolicy - [BackendSecurityPolicySpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicyspec) - [BackendSecurityPolicyStatus](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicystatus) - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicytype) -- [GCPContextCachingSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpcontextcachingspec) +- [ContextCacheSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-contextcachespec) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpcredentialsfile) - [GCPOIDCExchangeToken](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpoidcexchangetoken) - [GCPServiceAccountImpersonationConfig](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpserviceaccountimpersonationconfig) @@ -853,10 +853,10 @@ AIServiceBackendSpec details the AIServiceBackend configuration. required="false" description="BodyMutation defines the mutation of HTTP request body JSON fields that will be applied to the request
before sending it to the backend." /> @@ -1382,29 +1382,33 @@ BackendSecurityPolicyType specifies the type of auth mechanism used to access a required="false" description="" /> -#### GCPContextCachingSpec +#### ContextCacheSpec **Appears in:** - [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-aiservicebackendspec) -GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. +ContextCacheSpec configures context caching for a backend. + +Setting this field enables context caching: cache_control markers in requests are +resolved against the shared store at URL, and a cache entry is created when absent. +Leaving it unset disables caching entirely, and markers are ignored. ##### Fields @@ -3182,9 +3186,9 @@ MCPRouteList contains a list of MCPRoute. - [BackendSecurityPolicySpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicyspec) - [BackendSecurityPolicyStatus](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicystatus) - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicytype) +- [ContextCacheSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-contextcachespec) - [CredentialOverrideFromDynamicMetadata](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromdynamicmetadata) - [CredentialOverrideFromRequestHeaders](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromrequestheaders) -- [GCPContextCachingSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpcontextcachingspec) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpcredentialsfile) - [GCPOIDCExchangeToken](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpoidcexchangetoken) - [GCPServiceAccountImpersonationConfig](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpserviceaccountimpersonationconfig) @@ -3460,10 +3464,10 @@ AIServiceBackendSpec details the AIServiceBackend configuration. required="false" description="HeaderValueFilters filter individual values out of multi-valued request headers before sending
the request to this backend, so that one value an upstream provider does not accept does not
fail the whole request. At most one filter per header name.
Only GCPAnthropic and AWSAnthropic backends honor this, and today only for the
`anthropic-beta` header — a filter naming any other header is accepted but has no effect.
Setting it on a backend with any other schema is rejected." /> @@ -4027,6 +4031,36 @@ BackendSecurityPolicyType specifies the type of auth mechanism used to access a required="false" description="" /> +#### ContextCacheSpec + + + +**Appears in:** +- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-aiservicebackendspec) + +ContextCacheSpec configures context caching for a backend. + +Setting this field enables context caching: cache_control markers in requests are +resolved against the shared store at URL, and a cache entry is created when absent. +Leaving it unset disables caching entirely, and markers are ignored. + +##### Fields + + + + + + #### CredentialOverrideFromDynamicMetadata @@ -4087,32 +4121,6 @@ CredentialOverrideFromRequestHeaders sources the per-request credential from a r /> -#### GCPContextCachingSpec - - - -**Appears in:** -- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-aiservicebackendspec) - -GCPContextCachingSpec configures in-process Gemini context caching for a GCP Vertex AI backend. - -##### Fields - - - - - - #### GCPCredentialsFile From b33c7abf1de6b50d01697ee00a47a3938bc87083 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 30 Sep 2026 15:30:06 +0530 Subject: [PATCH 08/17] gcpcache: drop the create lock and the request-path list Creation was arbitrated with a Redis SET NX PX lock: one replica created the cache while the others polled for up to 5s inside ProcessRequestHeaders. The lock was not a guarantee, since a waiter that timed out went on to list and create anyway. It cost a blocking wait on the request path, up to ~100 Redis round trips per waiting request, and timeout constants that were never measured. Resolution on a store miss is now: create the cache, publish it to the store, and return it. There is no waiting on other replicas and no cachedContents.list call. Replicas that miss the same cold prefix at the same moment each create a cache and each keep their own. This is an accepted cost. Once any replica publishes a name, all replicas reuse it. Every create now reports Created with its real TokenCount. Previously the post-create re-list could switch to another replica's cache and report Created=false, which dropped a billed write from the cache-write token metrics. Duplicate creates now show up there. A create response without an expireTime is now an error. Its store TTL would be negative, so the write was skipped silently and every later request for that prefix created again. A leftover lock sentinel from an older replica fails to decode and is treated as a miss, so rolling deploys need no migration. The list code is removed from the request path. A background reconciler that repairs drift between Redis and GCP will follow separately. Signed-off-by: Sivanantham Chinnaiyan --- internal/gcpcache/gcpcache.go | 201 ++++--------------------- internal/gcpcache/gcpcache_test.go | 116 ++++++++++----- internal/gcpcache/redis.go | 97 +----------- internal/gcpcache/redis_test.go | 230 ++++++++++------------------- 4 files changed, 183 insertions(+), 461 deletions(-) diff --git a/internal/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go index 695b55aa52..6d5775511d 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/gcpcache/gcpcache.go @@ -11,12 +11,13 @@ // 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, lists Google cachedContents for the target region and model. -// 5. Creates the cache entry if not found, then returns the cache resource name and the -// non-cached messages. +// 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. // -// Concurrent resolutions of the same key are deduplicated by the store's create lock, -// which coordinates across gateway replicas rather than only within one process. +// 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. // // The CacheResolver interface allows the implementation to be replaced with an external // cache service in the future without changing the request-path callers. @@ -86,9 +87,10 @@ type CacheResolver interface { type resolver struct { httpClient *http.Client - // store holds resolved cache names. It is shared across replicas when backed by - // Redis, which is what keeps two replicas from creating the same cache. Its - // failures are never fatal: they are logged and treated as a miss. + // 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 CacheStore // logger records store failures, which are otherwise invisible because the request @@ -157,9 +159,8 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion }, nil } - // Store miss — resolve against GCP. Concurrent resolutions of the same key are - // arbitrated by the store's create lock (see resolveUncached), which deduplicates - // across replicas rather than only within this one. + // Store miss — create against GCP. See resolveUncached for how concurrent + // resolutions of the same key behave. result, err := r.resolveUncached(ctx, openAIReq, gcpAuth, cacheKey, ttl, contents, systemInstruction, geminiTools) if err != nil { return nil, err @@ -171,14 +172,14 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion return result, nil } -// resolveUncached performs the GCP-side resolution for a cache key: list, create if -// absent, then re-list to converge duplicate-create races. It deals only in the cached -// prefix, so it leaves ResolveResult.Messages unset for the caller to fill in per-request. +// resolveUncached creates a Google 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. // -// Before creating, it claims a store-level lock so that replicas racing on a cold prefix -// produce one create rather than N. Losing the race means waiting for the winner's result; -// the wait degrading into a plain resolution is safe, because the list below is the actual -// source of truth. +// 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, @@ -191,7 +192,6 @@ func (r *resolver) resolveUncached( region := gcpAuth.GCPRegion() project := gcpAuth.GCPProject() - // Get access token. tokenSrc := gcpAuth.GCPTokenSource() token, err := tokenSrc.Token() if err != nil { @@ -199,57 +199,15 @@ func (r *resolver) resolveUncached( } accessToken := token.AccessToken - // Claim the right to create this key. A replica that loses waits for the winner's - // result; if the winner never publishes one, it falls through and resolves itself. - locked := r.tryLock(ctx, cacheKey) - if !locked { - if e, ok := r.awaitLeader(ctx, cacheKey); ok { - return &ResolveResult{CacheName: e.cacheName, ExpireTime: e.expireTime}, nil - } - } - - // List existing caches and match by displayName (cacheKey). + // 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 Google caches. baseURL := fmt.Sprintf(gcpCachedContentsBasePath, region, project, region) - existingName, expireTime, err := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) - if err != nil { - if locked { - r.unlock(ctx, cacheKey) - } - return nil, fmt.Errorf("gcpcache: failed to list cached contents: %w", err) - } - - if existingName != "" { - r.storeSet(ctx, cacheKey, entry{cacheName: existingName, expireTime: expireTime}) - return &ResolveResult{ - CacheName: existingName, - ExpireTime: expireTime, - }, nil - } - - // Cache not found — create it. A replica that did not win the lock still creates - // here, having already waited for the winner without result; this is the duplicate - // the lock narrows but cannot fully eliminate. created, tokenCount, expireTime, err := r.createCache(ctx, baseURL, accessToken, openAIReq.Model, region, project, cacheKey, contents, systemInstruction, geminiTools, ttl) if err != nil { - // Release the claim so waiters retry immediately rather than blocking for the - // remainder of the lock TTL on a create that will never publish. - if locked { - r.unlock(ctx, cacheKey) - } return nil, fmt.Errorf("gcpcache: failed to create cached content: %w", err) } - // After create, re-list and prefer the oldest match to converge duplicate-create races. - finalName, finalExpire, listErr := r.listAndMatch(ctx, baseURL, accessToken, cacheKey, openAIReq.Model) - if listErr == nil && finalName != "" && finalName != created { - // Another replica created a cache with the same key; use the one from the list. - r.storeSet(ctx, cacheKey, entry{cacheName: finalName, expireTime: finalExpire}) - return &ResolveResult{ - CacheName: finalName, - ExpireTime: finalExpire, - }, nil - } - r.storeSet(ctx, cacheKey, entry{cacheName: created, expireTime: expireTime}) return &ResolveResult{ CacheName: created, @@ -303,52 +261,6 @@ func (r *resolver) storeSet(ctx context.Context, key string, e entry) { } } -// locker is implemented by stores that can arbitrate cache creation across replicas. -// Stores that cannot (such as the no-op store) simply never grant a lock, which leaves -// every caller resolving independently — correct, just without the deduplication. -type locker interface { - tryLock(ctx context.Context, key string) (bool, error) - unlock(ctx context.Context, key string) - awaitLeader(ctx context.Context, key string) (entry, bool, error) -} - -// tryLock claims the right to create key. It reports false when the store cannot lock, -// when the lock is held elsewhere, or on failure — all cases where the caller should -// resolve against Google rather than wait. -func (r *resolver) tryLock(ctx context.Context, key string) bool { - l, ok := r.store.(locker) - if !ok { - return false - } - won, err := l.tryLock(ctx, key) - if err != nil { - r.logger.Warn("gcpcache: cache store lock failed, proceeding without coordination", - slog.String("error", err.Error())) - return false - } - return won -} - -func (r *resolver) unlock(ctx context.Context, key string) { - if l, ok := r.store.(locker); ok { - l.unlock(ctx, key) - } -} - -// awaitLeader waits for the replica holding the lock to publish its result. -func (r *resolver) awaitLeader(ctx context.Context, key string) (entry, bool) { - l, ok := r.store.(locker) - if !ok { - return entry{}, false - } - e, found, err := l.awaitLeader(ctx, key) - if err != nil { - // Includes the timeout case: fall through and resolve rather than fail. - return entry{}, false - } - return e, found -} - // ----------------------------------------------------------------------- // Breakpoint and TTL helpers // ----------------------------------------------------------------------- @@ -504,70 +416,6 @@ func computeCacheKey(model string, contents []genai.Content, systemInstruction * // Google cachedContents REST API calls // ----------------------------------------------------------------------- -// cachedContentItem is the subset of the cachedContents list item we care about. -type cachedContentItem struct { - Name string `json:"name"` - DisplayName string `json:"displayName"` - ExpireTime string `json:"expireTime"` // RFC 3339 - Model string `json:"model"` -} - -type listResponse struct { - CachedContents []cachedContentItem `json:"cachedContents"` -} - -// listAndMatch lists cachedContents for the given region+project and returns the -// name of the first entry whose displayName matches cacheKey and whose model -// suffix matches the requested model. Returns ("", zero, nil) when not found. -func (r *resolver) listAndMatch(ctx context.Context, baseURL, accessToken, cacheKey, model string) (string, time.Time, error) { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL, nil) - if err != nil { - return "", time.Time{}, err - } - req.Header.Set("Authorization", "Bearer "+accessToken) - - resp, err := r.httpClient.Do(req) - if err != nil { - return "", time.Time{}, err - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", time.Time{}, err - } - if resp.StatusCode != http.StatusOK { - return "", time.Time{}, fmt.Errorf("list cachedContents returned HTTP %d: %s", resp.StatusCode, body) - } - - var lr listResponse - if err = json.Unmarshal(body, &lr); err != nil { - return "", time.Time{}, fmt.Errorf("failed to decode list response: %w", err) - } - - for _, item := range lr.CachedContents { - if item.DisplayName != cacheKey { - continue - } - // Model in the item is the full resource name; check the suffix. - if !modelMatchesSuffix(item.Model, model) { - continue - } - expireTime, _ := time.Parse(time.RFC3339, item.ExpireTime) - return item.Name, expireTime, nil - } - return "", time.Time{}, nil -} - -// modelMatchesSuffix checks whether the full model resource name ends with the -// requested short model name (e.g. "publishers/google/models/gemini-1.5-pro"). -func modelMatchesSuffix(fullModel, shortModel string) bool { - if fullModel == shortModel { - return true - } - suffix := "models/" + shortModel - return len(fullModel) >= len(suffix) && fullModel[len(fullModel)-len(suffix):] == suffix -} - // createCache creates a cached content in GCP and returns the new cache's resource name, // token count, and expiry time. func (r *resolver) createCache( @@ -618,6 +466,11 @@ func (r *resolver) createCache( 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) diff --git a/internal/gcpcache/gcpcache_test.go b/internal/gcpcache/gcpcache_test.go index e0a0141f1f..6e54d3b2ab 100644 --- a/internal/gcpcache/gcpcache_test.go +++ b/internal/gcpcache/gcpcache_test.go @@ -158,9 +158,7 @@ func resolverWithServer(srvURL string, store ...CacheStore) *resolver { }, s, nil).(*resolver) } -// memStore is an in-process CacheStore for tests. It implements only the CacheStore -// contract, not locker, so resolvers using it never coordinate — which is what the -// no-op-store path does in production too. +// memStore is an in-process CacheStore for tests. type memStore struct { mu sync.Mutex entries map[string]entry @@ -384,12 +382,15 @@ func TestResolver_CacheMiss_Creates(t *testing.T) { 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) - // list → create → re-list (post-create convergence check) - assert.Equal(t, 2, fake.lists()) + // The request path creates directly; it never lists. + assert.Equal(t, 0, fake.lists()) assert.Equal(t, 1, fake.creates()) } -func TestResolver_CacheHit_FromGoogleList(t *testing.T) { +// 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 reconciler's job. +func TestResolver_RequestPath_NeverLists(t *testing.T) { req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -412,16 +413,18 @@ func TestResolver_CacheHit_FromGoogleList(t *testing.T) { }) require.NoError(t, err) - fake := newFakeCacheServer(t, string(listBody), "", http.StatusOK) + createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `"}` + fake := newFakeCacheServer(t, string(listBody), createResp, http.StatusOK) r := resolverWithServer(fake.srv.URL) auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} res, err := r.Resolve(context.Background(), req, auth) require.NoError(t, err) require.NotNil(t, res) - assert.Equal(t, "projects/p/locations/us-central1/cachedContents/existing", res.CacheName) - assert.False(t, res.Created) - assert.Equal(t, 0, fake.creates()) + 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) { @@ -460,17 +463,8 @@ func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { key := computeKeyForRequest(t, req) expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) - listBody, _ := json.Marshal(map[string]interface{}{ - "cachedContents": []map[string]interface{}{ - { - "name": "projects/p/locations/us-central1/cachedContents/refreshed", - "displayName": key, - "model": "publishers/google/models/gemini-1.5-pro", - "expireTime": expireISO, - }, - }, - }) - fake := newFakeCacheServer(t, string(listBody), "", http.StatusOK) + 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. @@ -481,7 +475,55 @@ func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { 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.lists(), "a near-expiry store entry must trigger a list call") + 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.cacheName) +} + +// 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) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + + res, err := r.Resolve(context.Background(), req, auth) + 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) + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + req := &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("sys", ephemeralFields()), + userMsg("q"), + }, + } + + res, err := r.Resolve(context.Background(), req, auth) + require.Error(t, err) + assert.Nil(t, res) } func TestResolver_CreateFailure_ReturnsError(t *testing.T) { @@ -576,7 +618,10 @@ func TestResolver_TTLOverride_SentToGoogle(t *testing.T) { assert.Equal(t, "3600s", body.TTL) } -func TestResolver_DuplicateCreateRace_UsesListResult(t *testing.T) { +// 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{ @@ -587,11 +632,7 @@ func TestResolver_DuplicateCreateRace_UsesListResult(t *testing.T) { key := computeKeyForRequest(t, req) expireISO := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339) - listCalls := 0 - // First list: empty (cache not yet created by this replica). - // Create: returns "new-by-us". - // Second list (post-create): returns "older-by-other" — simulates another replica winning the race. - listBodyEmpty := `{"cachedContents":[]}` + var listCalls atomic.Int64 listBodyOther, _ := json.Marshal(map[string]interface{}{ "cachedContents": []map[string]interface{}{ { @@ -607,16 +648,12 @@ func TestResolver_DuplicateCreateRace_UsesListResult(t *testing.T) { w.Header().Set("Content-Type", "application/json") switch r.Method { case http.MethodGet: - listCalls++ + listCalls.Add(1) w.WriteHeader(http.StatusOK) - if listCalls == 1 { - _, _ = w.Write([]byte(listBodyEmpty)) - } else { - _, _ = w.Write(listBodyOther) - } + _, _ = 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 + `"}`)) + _, _ = w.Write([]byte(`{"name":"projects/p/locations/us-central1/cachedContents/new-by-us","expireTime":"` + expireISO + `","usageMetadata":{"totalTokenCount":512}}`)) } })) defer srv.Close() @@ -626,8 +663,9 @@ func TestResolver_DuplicateCreateRace_UsesListResult(t *testing.T) { res, err := r.Resolve(context.Background(), req, auth) require.NoError(t, err) require.NotNil(t, res) - // Resolver should prefer the entry from the post-create re-list (the other replica's entry). - assert.Equal(t, "projects/p/locations/us-central1/cachedContents/older-by-other", res.CacheName) - // Created should be false when we converged to another replica's cache. - assert.False(t, res.Created) + 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") } diff --git a/internal/gcpcache/redis.go b/internal/gcpcache/redis.go index 12b521f50e..3baf999e93 100644 --- a/internal/gcpcache/redis.go +++ b/internal/gcpcache/redis.go @@ -15,36 +15,7 @@ import ( "github.com/redis/go-redis/v9" ) -const ( - // lockSentinel marks a key as claimed by a replica that is creating the cache. - // A reader seeing it knows a create is in flight rather than that a name is stored. - lockSentinel = "\x00creating" - - // lockTTL bounds how long a create lock is held. It must exceed the p99 duration of - // a Google cachedContents create, or a second replica will claim the lock while the - // first is still working and both will create. It is deliberately not configurable: - // it is internal to the lock protocol, and widening the config surface later is - // additive whereas removing a field is a breaking change. - lockTTL = 30 * time.Second - - // lockWaitTime bounds how long a replica polls for the leader's result before giving - // up and resolving against Google itself. It must stay well under Envoy's ext_proc - // message timeout: a waiting request is blocked inside ProcessRequestHeaders. - lockWaitTime = 5 * time.Second - - // lockPollInterval is how often a waiter re-reads the key while the leader works. - lockPollInterval = 50 * time.Millisecond -) - -// errLockHeld reports that another replica holds the create lock and did not publish a -// result before lockWaitTime elapsed. -var errLockHeld = errors.New("gcpcache: create lock held by another replica") - // redisStore is a CacheStore backed by Redis, shared across gateway replicas. -// -// Beyond plain get/set it arbitrates cache creation: the first replica to claim a key -// creates the Google cache while the others wait for its result, so a cold prefix -// arriving at N replicas produces one create rather than N. type redisStore struct { client redis.UniversalClient } @@ -77,8 +48,7 @@ func parseRedisURL(url string) (*redis.Options, error) { return &redis.Options{Addr: url}, nil } -// Get implements CacheStore. A key holding the create sentinel is reported as a miss, -// since no cache name is available yet. +// Get implements CacheStore. func (s *redisStore) Get(ctx context.Context, key string) (entry, bool, error) { v, err := s.client.Get(ctx, key).Result() if err != nil { @@ -87,9 +57,6 @@ func (s *redisStore) Get(ctx context.Context, key string) (entry, bool, error) { } return entry{}, false, fmt.Errorf("gcpcache: redis get: %w", err) } - if v == lockSentinel { - return entry{}, false, nil - } e, err := decodeEntry(v) if err != nil { // A malformed value is treated as a miss rather than an error: it is not worth @@ -107,68 +74,6 @@ func (s *redisStore) Set(ctx context.Context, key string, e entry, ttl time.Dura return nil } -// tryLock attempts to claim key for creation. It reports whether the caller won. -// -// The winner must publish its result with Set (or release the claim with unlock) so that -// waiters are not stuck until the lock expires. -func (s *redisStore) tryLock(ctx context.Context, key string) (bool, error) { - won, err := s.client.SetNX(ctx, key, lockSentinel, lockTTL).Result() - if err != nil { - return false, fmt.Errorf("gcpcache: redis lock: %w", err) - } - return won, nil -} - -// unlock releases a claim that did not produce a cache name, so waiters retry promptly -// instead of blocking for the remainder of lockTTL. -// -// It deletes only a key still holding the sentinel: if a create succeeded and published a -// real name, or the lock expired and another replica re-claimed it, the key is left alone. -func (s *redisStore) unlock(ctx context.Context, key string) { - const script = `if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) end return 0` - // Best-effort: a failure here only means waiters fall back to Google after lockWaitTime. - _ = s.client.Eval(ctx, script, []string{key}, lockSentinel).Err() -} - -// awaitLeader polls key until the leader publishes a cache name, the claim disappears, or -// lockWaitTime elapses. -// -// A timeout returns errLockHeld, which the caller treats as "resolve against Google -// yourself". That degrades a stuck leader into an extra list rather than a duplicate -// create, since the resolver lists before it creates. -func (s *redisStore) awaitLeader(ctx context.Context, key string) (entry, bool, error) { - deadline := time.NewTimer(lockWaitTime) - defer deadline.Stop() - ticker := time.NewTicker(lockPollInterval) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - return entry{}, false, ctx.Err() - case <-deadline.C: - return entry{}, false, errLockHeld - case <-ticker.C: - v, err := s.client.Get(ctx, key).Result() - if err != nil { - if errors.Is(err, redis.Nil) { - // The leader released without publishing; caller should resolve. - return entry{}, false, errLockHeld - } - return entry{}, false, fmt.Errorf("gcpcache: redis get: %w", err) - } - if v == lockSentinel { - continue // Still creating. - } - e, decErr := decodeEntry(v) - if decErr != nil { - return entry{}, false, errLockHeld - } - return e, true, 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 entry) string { diff --git a/internal/gcpcache/redis_test.go b/internal/gcpcache/redis_test.go index f81ef3b135..80ec91932d 100644 --- a/internal/gcpcache/redis_test.go +++ b/internal/gcpcache/redis_test.go @@ -7,8 +7,11 @@ package gcpcache import ( "context" + "fmt" "net/http" + "net/http/httptest" "sync" + "sync/atomic" "testing" "time" @@ -71,21 +74,6 @@ func TestRedisStore_GetMiss(t *testing.T) { assert.False(t, ok) } -// A key holding the create sentinel is a claim, not a result: reporting it as a hit would -// hand the caller the sentinel string as a cache name. -func TestRedisStore_SentinelIsAMiss(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - - _, ok, err := s.Get(ctx, "k") - require.NoError(t, err) - assert.False(t, ok, "an in-flight create must not read as a cache hit") -} - // 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) { @@ -97,112 +85,6 @@ func TestRedisStore_MalformedValueIsAMiss(t *testing.T) { assert.False(t, ok) } -func TestRedisStore_TryLockIsExclusive(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - - first, err := s.tryLock(ctx, "k") - require.NoError(t, err) - assert.True(t, first) - - second, err := s.tryLock(ctx, "k") - require.NoError(t, err) - assert.False(t, second, "a second claim on a held key must lose") -} - -// unlock must not delete a key that has moved on from the sentinel, or it would erase a -// published cache name (or another replica's fresh claim after the lock expired). -func TestRedisStore_UnlockOnlyRemovesOwnSentinel(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - expire := time.Now().Add(10 * time.Minute) - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - require.NoError(t, s.Set(ctx, "k", entry{cacheName: "published", expireTime: expire}, time.Minute)) - - s.unlock(ctx, "k") - - got, ok, err := s.Get(ctx, "k") - require.NoError(t, err) - require.True(t, ok, "unlock must not erase a published result") - assert.Equal(t, "published", got.cacheName) -} - -func TestRedisStore_UnlockReleasesClaim(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - s.unlock(ctx, "k") - - again, err := s.tryLock(ctx, "k") - require.NoError(t, err) - assert.True(t, again, "a released claim must be re-claimable immediately") -} - -func TestRedisStore_AwaitLeader_ReturnsPublishedResult(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - expire := time.Now().Add(10 * time.Minute).UTC().Truncate(time.Second) - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - - go func() { - time.Sleep(100 * time.Millisecond) - _ = s.Set(ctx, "k", entry{cacheName: "from-leader", expireTime: expire}, time.Minute) - }() - - got, ok, err := s.awaitLeader(ctx, "k") - require.NoError(t, err) - require.True(t, ok) - assert.Equal(t, "from-leader", got.cacheName) -} - -// A leader that releases without publishing must not strand its waiters for the rest of -// lockWaitTime; they are told to resolve themselves as soon as the claim disappears. -func TestRedisStore_AwaitLeader_ClaimReleasedWithoutResult(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx := context.Background() - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - - go func() { - time.Sleep(100 * time.Millisecond) - s.unlock(ctx, "k") - }() - - start := time.Now() - _, ok, err := s.awaitLeader(ctx, "k") - require.ErrorIs(t, err, errLockHeld) - assert.False(t, ok) - assert.Less(t, time.Since(start), lockWaitTime, "waiter must not block for the full wait window") -} - -func TestRedisStore_AwaitLeader_HonorsContextCancellation(t *testing.T) { - s, _ := newTestRedisStore(t) - ctx, cancel := context.WithCancel(context.Background()) - - won, err := s.tryLock(ctx, "k") - require.NoError(t, err) - require.True(t, won) - - go func() { - time.Sleep(100 * time.Millisecond) - cancel() - }() - - _, _, err = s.awaitLeader(ctx, "k") - require.ErrorIs(t, err, context.Canceled) -} - 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. @@ -223,10 +105,9 @@ func TestEncodeDecodeEntry(t *testing.T) { // Cross-replica behavior // ----------------------------------------------------------------------- -// Two resolvers stand in for two gateway replicas: separate processes, one shared -// Redis. Deduplicating a cold prefix across them is the reason the store's create -// lock exists. -func TestResolver_CrossReplica_SingleCreate(t *testing.T) { +// 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) @@ -238,48 +119,93 @@ func TestResolver_CrossReplica_SingleCreate(t *testing.T) { return resolverWithServer(fake.srv.URL, store) } a, b := newReplica(), newReplica() + auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - req := func() *openai.ChatCompletionRequest { - return &openai.ChatCompletionRequest{ - Model: "gemini-1.5-pro", - Messages: []openai.ChatCompletionMessageParamUnion{ - systemMsg("You are helpful.", ephemeralFields()), - userMsg("Hello"), - }, + resA, err := a.Resolve(context.Background(), crossReplicaRequest(), auth) + 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(), auth) + 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) auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - var start, done sync.WaitGroup - start.Add(1) - results := make([]*ResolveResult, 2) - errs := make([]error, 2) - for i, r := range []*resolver{a, b} { + var done sync.WaitGroup + results := make([]*ResolveResult, replicas) + errs := make([]error, replicas) + for i := range replicas { + store, err := NewRedisStore(mr.Addr()) + require.NoError(t, err) + r := resolverWithServer(srv.URL, store) done.Add(1) go func() { defer done.Done() - start.Wait() - results[i], errs[i] = r.Resolve(context.Background(), req(), auth) + results[i], errs[i] = r.Resolve(context.Background(), crossReplicaRequest(), auth) }() } - start.Done() + for range replicas { + <-arrived + } + close(release) done.Wait() - for i := range results { + 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, results[i], "replica %d", i) - assert.Equal(t, "projects/p/locations/us-central1/cachedContents/new", results[i].CacheName) + 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.Equal(t, 1, fake.creates(), "two replicas sharing one Redis must create exactly one cache") + assert.Len(t, names, replicas, "each replica must keep the cache it created") +} - // Only the replica that performed the write may bill for it. - created := 0 - for _, res := range results { - if res.Created { - created++ - } +func crossReplicaRequest() *openai.ChatCompletionRequest { + return &openai.ChatCompletionRequest{ + Model: "gemini-1.5-pro", + Messages: []openai.ChatCompletionMessageParamUnion{ + systemMsg("You are helpful.", ephemeralFields()), + userMsg("Hello"), + }, } - assert.Equal(t, 1, created, "exactly one replica may report Created") } // The failure policy: an unreachable store degrades caching, it does not fail requests. From 932bb4c0fcf064b2edf0d2f23eabe1a8a3869da5 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 30 Sep 2026 16:17:55 +0530 Subject: [PATCH 09/17] gcpcache: add a background reconciler for store drift The request path no longer lists cachedContents, so a Google cache that the store has forgotten (after a Redis flush or eviction, a dropped write, or a cache created before Redis was configured) is not found again. Add a reconciler that repairs this off the request path. Each round lists cachedContents for one project and region and writes every gateway-created entry to the store with SET NX and a TTL pinned to its expireTime. SET NX means it only fills gaps and never overwrites a name that a request just published. Entries that are not keyed by a cache key were created by some other client and are skipped, as are entries near or past expiry and entries whose expireTime does not parse. The list API has no server-side filter, so the walk reads every entry. It requests the maximum page size, follows nextPageToken, treats a repeated token as the end, and stops after a fixed number of pages. A per-project/region gate, taken with SET NX and the interval as TTL, limits listing to one replica per interval across the fleet without leader election. A missed round has no effect on correctness. The reconciler never serves requests and never returns an error; a failed round is logged and the next one retries. Stores that cannot back it, such as the no-op store, have no reconciler. Nothing starts the reconciler yet; that comes with the resolver lifecycle change. Signed-off-by: Sivanantham Chinnaiyan --- internal/gcpcache/reconciler.go | 247 +++++++++++++++++++++++++++ internal/gcpcache/reconciler_test.go | 246 ++++++++++++++++++++++++++ internal/gcpcache/redis.go | 20 +++ internal/gcpcache/store.go | 7 +- 4 files changed, 517 insertions(+), 3 deletions(-) create mode 100644 internal/gcpcache/reconciler.go create mode 100644 internal/gcpcache/reconciler_test.go diff --git a/internal/gcpcache/reconciler.go b/internal/gcpcache/reconciler.go new file mode 100644 index 0000000000..5f9a30f65d --- /dev/null +++ b/internal/gcpcache/reconciler.go @@ -0,0 +1,247 @@ +// 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/filterapi" + "github.com/envoyproxy/ai-gateway/internal/json" +) + +// The reconciler repairs drift between the store and Google: caches that exist in +// cachedContents but that the store does not know about, after a Redis flush or +// eviction, a dropped write, or a cache created before the store was configured. +// +// It never serves requests. If it is late, stuck, or stopped, requests still succeed; +// the only effect is that more of them create a cache instead of reusing one. +// +// Each round: +// 1. Acquires a per-project/region interval gate with SET NX, so across the fleet at +// most one replica lists per interval. A replica that loses the gate skips the round. +// 2. Walks cachedContents.list page by page. The API has no server-side filter, so +// every entry in the project and region is read. +// 3. Writes each gateway-created entry to the store with SET NX and a TTL pinned to its +// expireTime. SET NX means it only fills gaps and never overwrites a name that a +// request just published. + +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 + + // listMaxPages bounds one round's walk, so a misbehaving nextPageToken chain cannot + // loop forever. Entries past the bound are picked up by the request path creating. + listMaxPages = 10 + + // reconcileGatePrefix namespaces gate keys. Cache keys are 64 hex characters, so the + // prefix cannot collide with them. + reconcileGatePrefix = "gcpcache:reconcile:" +) + +// reconcileStore is implemented by stores that can back the reconciler. It is kept out of +// CacheStore so that the no-op store and test stores are unaffected; a store that does +// not implement it simply has no reconciler. +type reconcileStore interface { + // 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) +} + +// cachedContentItem is the subset of a cachedContents list item the reconciler 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"` +} + +// reconcileStats summarizes one round, for logging and tests. +type reconcileStats struct { + gated bool // Another replica held the gate; nothing was listed. + pages int + seen int + written int + skipped int // Entries not written: foreign, expired, or unparseable. +} + +// runReconciler runs reconcile rounds for gcpAuth's project and region until ctx is +// done. It returns immediately if the store cannot back a reconciler (e.g. the no-op +// store). It never returns an error: failures are logged and the next round retries. +func (r *resolver) runReconciler(ctx context.Context, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) { + rs, ok := r.store.(reconcileStore) + if !ok { + return + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + r.reconcileRound(ctx, rs, gcpAuth, interval) + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} + +// reconcileRound runs one round and logs its failure, if any. +func (r *resolver) reconcileRound(ctx context.Context, rs reconcileStore, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) { + stats, err := r.reconcileOnce(ctx, rs, gcpAuth, interval) + if err != nil { + if ctx.Err() == nil { + r.logger.Warn("gcpcache: reconcile round failed", + slog.String("project", gcpAuth.GCPProject()), + slog.String("region", gcpAuth.GCPRegion()), + slog.String("error", err.Error())) + } + return + } + if !stats.gated { + r.logger.Debug("gcpcache: reconcile round complete", + slog.Int("pages", stats.pages), slog.Int("seen", stats.seen), + slog.Int("written", stats.written), slog.Int("skipped", stats.skipped)) + } +} + +// reconcileOnce performs one round. An error ends the round early; entries written +// before it stay written. +func (r *resolver) reconcileOnce(ctx context.Context, rs reconcileStore, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) (reconcileStats, error) { + var stats reconcileStats + region, project := gcpAuth.GCPRegion(), gcpAuth.GCPProject() + + won, err := rs.acquireGate(ctx, reconcileGatePrefix+project+"/"+region, interval) + if err != nil { + return stats, fmt.Errorf("acquire reconcile gate: %w", err) + } + if !won { + stats.gated = true + return stats, nil + } + + token, err := gcpAuth.GCPTokenSource().Token() + if err != nil { + return stats, fmt.Errorf("get GCP access token: %w", err) + } + + base, err := url.Parse(fmt.Sprintf(gcpCachedContentsBasePath, region, project, region)) + if err != nil { + return stats, fmt.Errorf("parse cachedContents URL: %w", err) + } + + pageToken := "" + for stats.pages < listMaxPages { + lr, err := r.listPage(ctx, base, token.AccessToken, pageToken) + if err != nil { + return stats, err + } + stats.pages++ + for _, item := range lr.CachedContents { + stats.seen++ + if r.reconcileItem(ctx, rs, item) { + stats.written++ + } else { + stats.skipped++ + } + } + // An empty token ends the list. A repeated one would loop, so treat it as the end. + if lr.NextPageToken == "" || lr.NextPageToken == pageToken { + return stats, nil + } + pageToken = lr.NextPageToken + } + r.logger.Warn("gcpcache: reconcile stopped at page limit; later entries were not read", + slog.String("project", project), slog.String("region", region), + slog.Int("maxPages", listMaxPages)) + return stats, nil +} + +// reconcileItem writes one list entry to the store. It reports whether it wrote. +func (r *resolver) reconcileItem(ctx context.Context, rs reconcileStore, item cachedContentItem) bool { + // Only entries the gateway created are keyed by a cache key. Anything else in the + // project was created by some other client and is not ours to index. + if !isCacheKey(item.DisplayName) || item.Name == "" { + return false + } + expireTime, err := time.Parse(time.RFC3339, item.ExpireTime) + if err != nil { + r.logger.Warn("gcpcache: skipping cachedContents entry with unparseable expireTime", + slog.String("name", item.Name), slog.String("expireTime", item.ExpireTime)) + return false + } + // storeGet treats entries this close to expiry as a miss, so writing one is pointless. + ttl := time.Until(expireTime) + if ttl < staleThreshold { + return false + } + wrote, err := rs.setNX(ctx, item.DisplayName, entry{cacheName: item.Name, expireTime: expireTime}, ttl) + if err != nil { + r.logger.Warn("gcpcache: reconcile store write failed", + slog.String("name", item.Name), slog.String("error", err.Error())) + return false + } + return wrote +} + +// 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/gcpcache/reconciler_test.go b/internal/gcpcache/reconciler_test.go new file mode 100644 index 0000000000..70c786b414 --- /dev/null +++ b/internal/gcpcache/reconciler_test.go @@ -0,0 +1,246 @@ +// 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/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var reconcileAuth = &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} + +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("reconciler 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 TestReconcile_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, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.Equal(t, 2, stats.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.cacheName) + } +} + +func TestReconcile_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), entry{cacheName: "c/published", expireTime: exp}, time.Minute)) + r := resolverWithServer(srv.URL, store) + + stats, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.Equal(t, 0, stats.written) + e, _, _ := store.Get(context.Background(), testKey(1)) + assert.Equal(t, "c/published", e.cacheName, "reconciler must not clobber a request's write") +} + +func TestReconcile_PageLimitStopsEndlessChain(t *testing.T) { + var calls atomic.Int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + n := calls.Add(1) + _, _ = fmt.Fprintf(w, `{"cachedContents":[],"nextPageToken":"t%d"}`, n) + })) + t.Cleanup(srv.Close) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + stats, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.Equal(t, listMaxPages, stats.pages) + assert.Equal(t, int64(listMaxPages), calls.Load()) +} + +func TestReconcile_RepeatedTokenEndsWalk(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{ + "": `{"cachedContents":[],"nextPageToken":"same"}`, + "same": `{"cachedContents":[],"nextPageToken":"same"}`, + }) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + + _, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.Equal(t, int64(2), calls.Load()) +} + +func TestReconcile_GateSkipsSecondReplica(t *testing.T) { + srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, mr := newTestRedisStore(t) + second, err := NewRedisStore(mr.Addr()) + require.NoError(t, err) + a := resolverWithServer(srv.URL, store) + b := resolverWithServer(srv.URL, second) + + statsA, err := a.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.False(t, statsA.gated) + statsB, err := b.reconcileOnce(context.Background(), second.(*redisStore), reconcileAuth, time.Minute) + require.NoError(t, err) + assert.True(t, statsB.gated) + assert.Equal(t, int64(1), calls.Load()) + + // Once the interval passes, a replica may reconcile again. + mr.FastForward(time.Minute + time.Second) + statsB, err = b.reconcileOnce(context.Background(), second.(*redisStore), reconcileAuth, time.Minute) + require.NoError(t, err) + assert.False(t, statsB.gated) +} + +func TestReconcile_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 := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + stats, err := r.reconcileOnce(context.Background(), store, other, time.Minute) + 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 TestReconcile_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 := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.Equal(t, 4, stats.seen) + assert.Equal(t, 1, stats.written) + assert.Equal(t, 3, 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 TestReconcile_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 := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.NoError(t, err) + assert.InDelta(t, (5 * time.Minute).Seconds(), mr.TTL(testKey(1)).Seconds(), 5) +} + +func TestReconcile_ListFailureReturnsErrorWithoutPanicking(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 := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + require.ErrorContains(t, err, "HTTP 500") + // reconcileRound is the caller the loop uses; it must swallow the error. + r.reconcileRound(context.Background(), store, reconcileAuth, time.Minute) +} + +func TestRunReconciler_NoopStoreReturnsImmediately(t *testing.T) { + r := New(nil, nil, nil).(*resolver) + done := make(chan struct{}) + go func() { + r.runReconciler(context.Background(), reconcileAuth, time.Minute) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runReconciler must return at once for a store that cannot reconcile") + } +} + +func TestRunReconciler_StopsOnContextCancel(t *testing.T) { + srv, _, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) + store, _ := newTestRedisStore(t) + r := resolverWithServer(srv.URL, store) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + r.runReconciler(ctx, reconcileAuth, time.Hour) + close(done) + }() + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runReconciler must return when its context is cancelled") + } +} + +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])) +} diff --git a/internal/gcpcache/redis.go b/internal/gcpcache/redis.go index 3baf999e93..aebc5cf7c6 100644 --- a/internal/gcpcache/redis.go +++ b/internal/gcpcache/redis.go @@ -74,6 +74,26 @@ func (s *redisStore) Set(ctx context.Context, key string, e entry, ttl time.Dura return nil } +// setNX implements reconcileStore. +func (s *redisStore) setNX(ctx context.Context, key string, e entry, ttl time.Duration) (bool, error) { + wrote, err := s.client.SetNX(ctx, key, encodeEntry(e), ttl).Result() + if err != nil { + return false, fmt.Errorf("gcpcache: redis setnx: %w", err) + } + return wrote, nil +} + +// acquireGate implements reconcileStore. +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("gcpcache: redis gate: %w", err) + } + return won, nil +} + +var _ reconcileStore = (*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 entry) string { diff --git a/internal/gcpcache/store.go b/internal/gcpcache/store.go index e643389b2f..bd422e0adf 100644 --- a/internal/gcpcache/store.go +++ b/internal/gcpcache/store.go @@ -19,9 +19,10 @@ type entry struct { // CacheStore is the storage layer behind the resolver, mapping a deterministic cache // key to the Google cache name it resolved to. // -// The store is shared across gateway replicas, which is what prevents two replicas -// from independently creating a cache for the same prefix. Implementations must be -// safe for concurrent use. +// The store is shared across gateway replicas, so a cache name published by one +// replica is reused by all of them. It is a cache, not a lock: replicas that miss the +// same cold prefix at the same moment may each create a Google cache, and the resolver +// tolerates that. Implementations must be safe for concurrent use. // // Store errors are never fatal to a request: the resolver treats them as a cache miss // and proceeds to Google. See the failure policy in the Redis cache store proposal. From f8842729cbc4a7f19ee65b3fc728ed15ef482f66 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 30 Sep 2026 17:10:11 +0530 Subject: [PATCH 10/17] contextcache: move the store and reconcile loop out of gcpcache The shared store and the reconcile loop do not depend on any provider API, but they lived in gcpcache. Move them into a provider-neutral internal/contextcache package so a second provider only has to supply its own API calls. contextcache now holds Store, Entry, NoopStore, the Redis store, and a ReconcileStore interface with SetNX and AcquireGate for stores that can back a reconciler. It also holds Reconciler, which runs the ticker, the per-scope gate, and the SET NX writes, and StaleThreshold. A provider plugs in through Source: GateKey names the scope's gate, and List reports every gateway-created cache in the scope with its cache key. gcpcache keeps the resolver, the cache key, createCache, and a Source that walks cachedContents.list with pagination and the page limit. That source now drops caches created by other clients, and caches whose expireTime does not parse, before the reconciler sees them, so they no longer count as seen or skipped in round stats. Redis keys, the stored value format, and the gcpcache:reconcile: gate prefix are unchanged, so old and new replicas interoperate during a rolling deploy. Store errors now start with "contextcache:" instead of "gcpcache:". Signed-off-by: Sivanantham Chinnaiyan --- internal/contextcache/reconciler.go | 139 +++++++++++++++ internal/contextcache/reconciler_test.go | 124 ++++++++++++++ internal/{gcpcache => contextcache}/redis.go | 58 +++---- internal/contextcache/redis_test.go | 95 +++++++++++ internal/contextcache/store.go | 70 ++++++++ internal/extproc/server.go | 3 +- internal/gcpcache/gcpcache.go | 32 ++-- internal/gcpcache/gcpcache_test.go | 21 +-- internal/gcpcache/reconciler.go | 169 ++++++------------- internal/gcpcache/reconciler_test.go | 114 ++++++------- internal/gcpcache/redis_test.go | 86 +--------- internal/gcpcache/store.go | 49 ------ 12 files changed, 594 insertions(+), 366 deletions(-) create mode 100644 internal/contextcache/reconciler.go create mode 100644 internal/contextcache/reconciler_test.go rename internal/{gcpcache => contextcache}/redis.go (56%) create mode 100644 internal/contextcache/redis_test.go create mode 100644 internal/contextcache/store.go delete mode 100644 internal/gcpcache/store.go diff --git a/internal/contextcache/reconciler.go b/internal/contextcache/reconciler.go new file mode 100644 index 0000000000..8183e71081 --- /dev/null +++ b/internal/contextcache/reconciler.go @@ -0,0 +1,139 @@ +// 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" + "log/slog" + "time" +) + +// 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 reconciler 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 reconciliation. +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 reconcile 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 +} + +// Reconciler 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. +// +// It never serves requests. If it is late, stuck, or stopped, requests still succeed; +// the only effect is that more of them create a cache instead of reusing one. +// +// 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 Reconciler struct { + Store ReconcileStore + 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 *Reconciler) 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 *Reconciler) round(ctx context.Context) { + stats, err := rc.RunOnce(ctx) + if err != nil { + if ctx.Err() == nil { + rc.logger().Warn("contextcache: reconcile round failed", + slog.String("gate", rc.Source.GateKey()), slog.String("error", err.Error())) + } + return + } + if !stats.Gated { + rc.logger().Debug("contextcache: reconcile 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 *Reconciler) 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 reconcile 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 *Reconciler) 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: reconcile store write failed", + slog.String("name", e.Name), slog.String("error", err.Error())) + return false + } + return wrote +} + +func (rc *Reconciler) logger() *slog.Logger { + if rc.Logger == nil { + return slog.New(slog.DiscardHandler) + } + return rc.Logger +} diff --git a/internal/contextcache/reconciler_test.go b/internal/contextcache/reconciler_test.go new file mode 100644 index 0000000000..afa6ae69b9 --- /dev/null +++ b/internal/contextcache/reconciler_test.go @@ -0,0 +1,124 @@ +// 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" + "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 +} + +func newTestReconciler(t *testing.T, src *fakeSource) (*Reconciler, ReconcileStore) { + t.Helper() + store, _ := newTestRedisStore(t) + return &Reconciler{Store: store, Source: src, Interval: time.Minute}, store +} + +func TestReconciler_WritesEntries(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestReconciler(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 TestReconciler_DoesNotOverwrite(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestReconciler(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 TestReconciler_SkipsNearExpiry(t *testing.T) { + rc, store := newTestReconciler(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 TestReconciler_GateSkipsSecondRound(t *testing.T) { + src := &fakeSource{gate: "g"} + rc, _ := newTestReconciler(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 TestReconciler_ListErrorKeepsEarlierWrites(t *testing.T) { + exp := time.Now().Add(5 * time.Minute) + rc, store := newTestReconciler(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 TestReconciler_RunStopsOnCancel(t *testing.T) { + rc, _ := newTestReconciler(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") + } +} diff --git a/internal/gcpcache/redis.go b/internal/contextcache/redis.go similarity index 56% rename from internal/gcpcache/redis.go rename to internal/contextcache/redis.go index aebc5cf7c6..b62052336e 100644 --- a/internal/gcpcache/redis.go +++ b/internal/contextcache/redis.go @@ -3,7 +3,7 @@ // The full text of the Apache license is available in the LICENSE file at // the root of the repo. -package gcpcache +package contextcache import ( "context" @@ -15,17 +15,17 @@ import ( "github.com/redis/go-redis/v9" ) -// redisStore is a CacheStore backed by Redis, shared across gateway replicas. +// redisStore is a Store backed by Redis, shared across gateway replicas. type redisStore struct { client redis.UniversalClient } -// NewRedisStore returns a CacheStore backed by the Redis instance at url, which accepts +// NewRedisStore returns a ReconcileStore 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 the resolver treats as a miss. -func NewRedisStore(url string) (CacheStore, error) { +// therefore surfaces as an error from Get/Set, which resolvers treat as a miss. +func NewRedisStore(url string) (ReconcileStore, error) { opts, err := parseRedisURL(url) if err != nil { return nil, err @@ -36,80 +36,78 @@ func NewRedisStore(url string) (CacheStore, error) { // parseRedisURL accepts both a scheme-qualified URL and a bare host:port. func parseRedisURL(url string) (*redis.Options, error) { if url == "" { - return nil, errors.New("gcpcache: redis url is empty") + return nil, errors.New("contextcache: redis url is empty") } if strings.Contains(url, "://") { opts, err := redis.ParseURL(url) if err != nil { - return nil, fmt.Errorf("gcpcache: invalid redis url: %w", err) + return nil, fmt.Errorf("contextcache: invalid redis url: %w", err) } return opts, nil } return &redis.Options{Addr: url}, nil } -// Get implements CacheStore. -func (s *redisStore) Get(ctx context.Context, key string) (entry, bool, error) { +// Get implements Store. +func (s *redisStore) Get(ctx context.Context, key string) (Entry, bool, error) { v, err := s.client.Get(ctx, key).Result() if err != nil { if errors.Is(err, redis.Nil) { - return entry{}, false, nil + return Entry{}, false, nil } - return entry{}, false, fmt.Errorf("gcpcache: redis get: %w", err) + return 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 entry{}, false, nil + return Entry{}, false, nil } return e, true, nil } -// Set implements CacheStore. -func (s *redisStore) Set(ctx context.Context, key string, e entry, ttl time.Duration) error { +// Set implements Store. +func (s *redisStore) Set(ctx context.Context, key string, e Entry, ttl time.Duration) error { if err := s.client.Set(ctx, key, encodeEntry(e), ttl).Err(); err != nil { - return fmt.Errorf("gcpcache: redis set: %w", err) + return fmt.Errorf("contextcache: redis set: %w", err) } return nil } -// setNX implements reconcileStore. -func (s *redisStore) setNX(ctx context.Context, key string, e entry, ttl time.Duration) (bool, error) { +// SetNX implements ReconcileStore. +func (s *redisStore) SetNX(ctx context.Context, key string, e Entry, ttl time.Duration) (bool, error) { wrote, err := s.client.SetNX(ctx, key, encodeEntry(e), ttl).Result() if err != nil { - return false, fmt.Errorf("gcpcache: redis setnx: %w", err) + return false, fmt.Errorf("contextcache: redis setnx: %w", err) } return wrote, nil } -// acquireGate implements reconcileStore. -func (s *redisStore) acquireGate(ctx context.Context, key string, ttl time.Duration) (bool, error) { +// AcquireGate implements ReconcileStore. +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("gcpcache: redis gate: %w", err) + return false, fmt.Errorf("contextcache: redis gate: %w", err) } return won, nil } -var _ reconcileStore = (*redisStore)(nil) +var _ ReconcileStore = (*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 entry) string { - return e.expireTime.UTC().Format(time.RFC3339) + "|" + e.cacheName +func encodeEntry(e Entry) string { + return e.ExpireTime.UTC().Format(time.RFC3339) + "|" + e.Name } -func decodeEntry(s string) (entry, error) { +func decodeEntry(s string) (Entry, error) { expiry, name, ok := strings.Cut(s, "|") if !ok || name == "" { - return entry{}, fmt.Errorf("gcpcache: malformed cache entry %q", s) + return Entry{}, fmt.Errorf("contextcache: malformed cache entry %q", s) } t, err := time.Parse(time.RFC3339, expiry) if err != nil { - return entry{}, fmt.Errorf("gcpcache: malformed cache entry expiry %q: %w", expiry, err) + return Entry{}, fmt.Errorf("contextcache: malformed cache entry expiry %q: %w", expiry, err) } - return entry{cacheName: name, expireTime: t}, nil + return Entry{Name: name, ExpireTime: t}, nil } - -var _ CacheStore = (*redisStore)(nil) diff --git a/internal/contextcache/redis_test.go b/internal/contextcache/redis_test.go new file mode 100644 index 0000000000..7030f58dbb --- /dev/null +++ b/internal/contextcache/redis_test.go @@ -0,0 +1,95 @@ +// 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" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// 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 := NewRedisStore(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", 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 := 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/store.go b/internal/contextcache/store.go new file mode 100644 index 0000000000..dd9d8db874 --- /dev/null +++ b/internal/contextcache/store.go @@ -0,0 +1,70 @@ +// 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 +// reconciler 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 reconciler. +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. It is a cache, not a lock: 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 +} + +// ReconcileStore is implemented by stores that can back a Reconciler. It is kept out of +// Store so that the no-op store and test stores need not implement it; a store that does +// not implement it simply has no reconciler. +type ReconcileStore interface { + Store + // 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. +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 } + +var _ Store = NoopStore{} diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 6a954ba184..ebf1318e00 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -26,6 +26,7 @@ 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/filterapi" "github.com/envoyproxy/ai-gateway/internal/gcpcache" "github.com/envoyproxy/ai-gateway/internal/internalapi" @@ -115,7 +116,7 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error key := rb.Backend.Name + "|" + cc.URL resolver, ok := s.cacheResolvers[key] if !ok { - store, err := gcpcache.NewRedisStore(cc.URL) + store, err := contextcache.NewRedisStore(cc.URL) if err != nil { // A bad URL disables caching for this backend rather than failing the // reload: the store is an optimization, and rejecting the whole config diff --git a/internal/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go index 6d5775511d..c396f478f8 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/gcpcache/gcpcache.go @@ -38,6 +38,7 @@ import ( "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" @@ -49,11 +50,6 @@ const ( // gcpCachedContentsBasePath is the base URL for the Vertex AI cachedContents REST API. gcpCachedContentsBasePath = "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/cachedContents" - - // staleThreshold is how close to its expiry a stored entry must be before it is - // treated as a miss, so there is time to re-resolve before the Google cache - // disappears underneath an in-flight request. - staleThreshold = 10 * time.Second ) // ResolveResult holds the outcome of a successful cache resolution. @@ -91,7 +87,7 @@ type resolver struct { // 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 CacheStore + store contextcache.Store // logger records store failures, which are otherwise invisible because the request // proceeds normally. @@ -103,12 +99,12 @@ type resolver struct { // httpClient is used for calls to the Google 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. logger may be nil. -func New(httpClient *http.Client, store CacheStore, logger *slog.Logger) CacheResolver { +func New(httpClient *http.Client, store contextcache.Store, logger *slog.Logger) CacheResolver { if httpClient == nil { httpClient = &http.Client{Timeout: 30 * time.Second} } if store == nil { - store = noopStore{} + store = contextcache.NoopStore{} } if logger == nil { logger = slog.New(slog.DiscardHandler) @@ -153,9 +149,9 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // request. Google-side failures below are a different matter and do fail fast. if e, ok := r.storeGet(ctx, cacheKey); ok { return &ResolveResult{ - CacheName: e.cacheName, + CacheName: e.Name, Messages: remainderMessages, - ExpireTime: e.expireTime, + ExpireTime: e.ExpireTime, }, nil } @@ -208,7 +204,7 @@ func (r *resolver) resolveUncached( return nil, fmt.Errorf("gcpcache: failed to create cached content: %w", err) } - r.storeSet(ctx, cacheKey, entry{cacheName: created, expireTime: expireTime}) + r.storeSet(ctx, cacheKey, contextcache.Entry{Name: created, ExpireTime: expireTime}) return &ResolveResult{ CacheName: created, Created: true, @@ -230,28 +226,28 @@ func (r *resolver) resolveUncached( // ----------------------------------------------------------------------- // storeGet reads a resolved cache name, reporting a miss on any failure. -func (r *resolver) storeGet(ctx context.Context, key string) (entry, bool) { +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 entry{}, false + return contextcache.Entry{}, false } if !ok { - return entry{}, false + 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) < staleThreshold { - return entry{}, false + if time.Until(e.ExpireTime) < contextcache.StaleThreshold { + return contextcache.Entry{}, false } return e, true } // storeSet records a resolved cache name, expiring it with the Google entry itself so // the store cannot outlive what it points at. -func (r *resolver) storeSet(ctx context.Context, key string, e entry) { - ttl := time.Until(e.expireTime) +func (r *resolver) storeSet(ctx context.Context, key string, e contextcache.Entry) { + ttl := time.Until(e.ExpireTime) if ttl <= 0 { return } diff --git a/internal/gcpcache/gcpcache_test.go b/internal/gcpcache/gcpcache_test.go index 6e54d3b2ab..f453f1604e 100644 --- a/internal/gcpcache/gcpcache_test.go +++ b/internal/gcpcache/gcpcache_test.go @@ -22,6 +22,7 @@ import ( "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" @@ -146,9 +147,9 @@ func (t *redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) // 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 ...CacheStore) *resolver { +func resolverWithServer(srvURL string, store ...contextcache.Store) *resolver { host := srvURL[len("http://"):] - var s CacheStore + var s contextcache.Store if len(store) > 0 { s = store[0] } @@ -158,10 +159,10 @@ func resolverWithServer(srvURL string, store ...CacheStore) *resolver { }, s, nil).(*resolver) } -// memStore is an in-process CacheStore for tests. +// memStore is an in-process contextcache.Store for tests. type memStore struct { mu sync.Mutex - entries map[string]entry + 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. @@ -170,12 +171,12 @@ type memStore struct { sets atomic.Int64 } -func newMemStore() *memStore { return &memStore{entries: map[string]entry{}} } +func newMemStore() *memStore { return &memStore{entries: map[string]contextcache.Entry{}} } -func (m *memStore) Get(_ context.Context, key string) (entry, bool, error) { +func (m *memStore) Get(_ context.Context, key string) (contextcache.Entry, bool, error) { m.gets.Add(1) if m.getErr != nil { - return entry{}, false, m.getErr + return contextcache.Entry{}, false, m.getErr } m.mu.Lock() defer m.mu.Unlock() @@ -183,7 +184,7 @@ func (m *memStore) Get(_ context.Context, key string) (entry, bool, error) { return e, ok, nil } -func (m *memStore) Set(_ context.Context, key string, e entry, _ time.Duration) error { +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 @@ -197,7 +198,7 @@ func (m *memStore) Set(_ context.Context, key string, e entry, _ time.Duration) func (m *memStore) seed(key, cacheName string, expireTime time.Time) { m.mu.Lock() defer m.mu.Unlock() - m.entries[key] = entry{cacheName: cacheName, expireTime: expireTime} + m.entries[key] = contextcache.Entry{Name: cacheName, ExpireTime: expireTime} } // computeKeyForRequest replicates the key generation for use in test assertions. @@ -482,7 +483,7 @@ func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { 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.cacheName) + 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), diff --git a/internal/gcpcache/reconciler.go b/internal/gcpcache/reconciler.go index 5f9a30f65d..b55f569b61 100644 --- a/internal/gcpcache/reconciler.go +++ b/internal/gcpcache/reconciler.go @@ -16,25 +16,14 @@ import ( "strconv" "time" + "github.com/envoyproxy/ai-gateway/internal/contextcache" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/json" ) -// The reconciler repairs drift between the store and Google: caches that exist in -// cachedContents but that the store does not know about, after a Redis flush or -// eviction, a dropped write, or a cache created before the store was configured. -// -// It never serves requests. If it is late, stuck, or stopped, requests still succeed; -// the only effect is that more of them create a cache instead of reusing one. -// -// Each round: -// 1. Acquires a per-project/region interval gate with SET NX, so across the fleet at -// most one replica lists per interval. A replica that loses the gate skips the round. -// 2. Walks cachedContents.list page by page. The API has no server-side filter, so -// every entry in the project and region is read. -// 3. Writes each gateway-created entry to the store with SET NX and a TTL pinned to its -// expireTime. SET NX means it only fills gaps and never overwrites a name that a -// request just published. +// This file supplies the GCP side of reconciliation: a contextcache.Source that walks +// cachedContents.list for one project and region. The loop, gate, and store writes live +// in contextcache.Reconciler. const ( // listPageSize is the documented maximum for cachedContents.list; larger values are @@ -46,20 +35,11 @@ const ( listMaxPages = 10 // reconcileGatePrefix namespaces gate keys. Cache keys are 64 hex characters, so the - // prefix cannot collide with them. + // prefix cannot collide with them. It is unchanged from before the store moved to + // contextcache, so old and new replicas share gates during a rolling deploy. reconcileGatePrefix = "gcpcache:reconcile:" ) -// reconcileStore is implemented by stores that can back the reconciler. It is kept out of -// CacheStore so that the no-op store and test stores are unaffected; a store that does -// not implement it simply has no reconciler. -type reconcileStore interface { - // 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) -} - // cachedContentItem is the subset of a cachedContents list item the reconciler uses. type cachedContentItem struct { Name string `json:"name"` @@ -72,131 +52,86 @@ type listResponse struct { NextPageToken string `json:"nextPageToken"` } -// reconcileStats summarizes one round, for logging and tests. -type reconcileStats struct { - gated bool // Another replica held the gate; nothing was listed. - pages int - seen int - written int - skipped int // Entries not written: foreign, expired, or unparseable. +// newReconciler returns a reconciler for gcpAuth's project and region. It reports false +// when the store cannot back one, such as the no-op store. +func (r *resolver) newReconciler(gcpAuth filterapi.GCPAuthHandler, interval time.Duration) (*contextcache.Reconciler, bool) { + rs, ok := r.store.(contextcache.ReconcileStore) + if !ok { + return nil, false + } + return &contextcache.Reconciler{ + Store: rs, + Source: &listSource{r: r, auth: gcpAuth}, + Interval: interval, + Logger: r.logger, + }, true } -// runReconciler runs reconcile rounds for gcpAuth's project and region until ctx is -// done. It returns immediately if the store cannot back a reconciler (e.g. the no-op -// store). It never returns an error: failures are logged and the next round retries. -func (r *resolver) runReconciler(ctx context.Context, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) { - rs, ok := r.store.(reconcileStore) - if !ok { - return - } - ticker := time.NewTicker(interval) - defer ticker.Stop() - for { - r.reconcileRound(ctx, rs, gcpAuth, interval) - select { - case <-ctx.Done(): - return - case <-ticker.C: - } - } +// 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 } -// reconcileRound runs one round and logs its failure, if any. -func (r *resolver) reconcileRound(ctx context.Context, rs reconcileStore, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) { - stats, err := r.reconcileOnce(ctx, rs, gcpAuth, interval) - if err != nil { - if ctx.Err() == nil { - r.logger.Warn("gcpcache: reconcile round failed", - slog.String("project", gcpAuth.GCPProject()), - slog.String("region", gcpAuth.GCPRegion()), - slog.String("error", err.Error())) - } - return - } - if !stats.gated { - r.logger.Debug("gcpcache: reconcile round complete", - slog.Int("pages", stats.pages), slog.Int("seen", stats.seen), - slog.Int("written", stats.written), slog.Int("skipped", stats.skipped)) - } +// GateKey implements contextcache.Source. +func (s *listSource) GateKey() string { + return reconcileGatePrefix + s.auth.GCPProject() + "/" + s.auth.GCPRegion() } -// reconcileOnce performs one round. An error ends the round early; entries written -// before it stay written. -func (r *resolver) reconcileOnce(ctx context.Context, rs reconcileStore, gcpAuth filterapi.GCPAuthHandler, interval time.Duration) (reconcileStats, error) { - var stats reconcileStats - region, project := gcpAuth.GCPRegion(), gcpAuth.GCPProject() +// 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 + region, project := s.auth.GCPRegion(), s.auth.GCPProject() - won, err := rs.acquireGate(ctx, reconcileGatePrefix+project+"/"+region, interval) + token, err := s.auth.GCPTokenSource().Token() if err != nil { - return stats, fmt.Errorf("acquire reconcile gate: %w", err) - } - if !won { - stats.gated = true - return stats, nil + return fmt.Errorf("get GCP access token: %w", err) } - - token, err := gcpAuth.GCPTokenSource().Token() - if err != nil { - return stats, fmt.Errorf("get GCP access token: %w", err) - } - base, err := url.Parse(fmt.Sprintf(gcpCachedContentsBasePath, region, project, region)) if err != nil { - return stats, fmt.Errorf("parse cachedContents URL: %w", err) + return fmt.Errorf("parse cachedContents URL: %w", err) } pageToken := "" - for stats.pages < listMaxPages { - lr, err := r.listPage(ctx, base, token.AccessToken, pageToken) + for s.pages < listMaxPages { + lr, err := s.r.listPage(ctx, base, token.AccessToken, pageToken) if err != nil { - return stats, err + return err } - stats.pages++ + s.pages++ for _, item := range lr.CachedContents { - stats.seen++ - if r.reconcileItem(ctx, rs, item) { - stats.written++ - } else { - stats.skipped++ + if key, e, ok := s.toEntry(item); ok { + add(key, e) } } // An empty token ends the list. A repeated one would loop, so treat it as the end. if lr.NextPageToken == "" || lr.NextPageToken == pageToken { - return stats, nil + return nil } pageToken = lr.NextPageToken } - r.logger.Warn("gcpcache: reconcile stopped at page limit; later entries were not read", + s.r.logger.Warn("gcpcache: reconcile stopped at page limit; later entries were not read", slog.String("project", project), slog.String("region", region), slog.Int("maxPages", listMaxPages)) - return stats, nil + return nil } -// reconcileItem writes one list entry to the store. It reports whether it wrote. -func (r *resolver) reconcileItem(ctx context.Context, rs reconcileStore, item cachedContentItem) bool { - // Only entries the gateway created are keyed by a cache key. Anything else in the - // project was created by some other client and is not ours to index. +// 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 false + return "", contextcache.Entry{}, false } expireTime, err := time.Parse(time.RFC3339, item.ExpireTime) if err != nil { - r.logger.Warn("gcpcache: skipping cachedContents entry with unparseable expireTime", + s.r.logger.Warn("gcpcache: skipping cachedContents entry with unparseable expireTime", slog.String("name", item.Name), slog.String("expireTime", item.ExpireTime)) - return false - } - // storeGet treats entries this close to expiry as a miss, so writing one is pointless. - ttl := time.Until(expireTime) - if ttl < staleThreshold { - return false - } - wrote, err := rs.setNX(ctx, item.DisplayName, entry{cacheName: item.Name, expireTime: expireTime}, ttl) - if err != nil { - r.logger.Warn("gcpcache: reconcile store write failed", - slog.String("name", item.Name), slog.String("error", err.Error())) - return false + return "", contextcache.Entry{}, false } - return wrote + 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 diff --git a/internal/gcpcache/reconciler_test.go b/internal/gcpcache/reconciler_test.go index 70c786b414..43cece4889 100644 --- a/internal/gcpcache/reconciler_test.go +++ b/internal/gcpcache/reconciler_test.go @@ -15,12 +15,34 @@ import ( "testing" "time" + "github.com/alicebob/miniredis/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" ) var reconcileAuth = &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.ReconcileStore, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + s, err := contextcache.NewRedisStore(mr.Addr()) + require.NoError(t, err) + return s, mr +} + +// reconcileOnce runs one round for auth against r's store and returns its stats and the +// number of list pages read. +func reconcileOnce(t *testing.T, r *resolver, auth *fakeGCPAuth) (contextcache.Stats, int, error) { + t.Helper() + rc, ok := r.newReconciler(auth, time.Minute) + require.True(t, ok) + 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 { @@ -60,10 +82,10 @@ func TestReconcile_MultiPageWalkWritesEveryEntry(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - stats, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + stats, pages, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) - assert.Equal(t, 2, stats.pages) - assert.Equal(t, 2, stats.written) + 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") @@ -72,7 +94,7 @@ func TestReconcile_MultiPageWalkWritesEveryEntry(t *testing.T) { 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.cacheName) + assert.Equal(t, name, e.Name) } } @@ -82,14 +104,14 @@ func TestReconcile_DoesNotOverwriteExistingEntry(t *testing.T) { "": `{"cachedContents":[` + itemJSON(testKey(1), "c/from-list", exp) + `]}`, }) store, _ := newTestRedisStore(t) - require.NoError(t, store.Set(context.Background(), testKey(1), entry{cacheName: "c/published", expireTime: exp}, time.Minute)) + 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 := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + stats, _, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) - assert.Equal(t, 0, stats.written) + assert.Equal(t, 0, stats.Written) e, _, _ := store.Get(context.Background(), testKey(1)) - assert.Equal(t, "c/published", e.cacheName, "reconciler must not clobber a request's write") + assert.Equal(t, "c/published", e.Name, "reconciler must not clobber a request's write") } func TestReconcile_PageLimitStopsEndlessChain(t *testing.T) { @@ -102,9 +124,9 @@ func TestReconcile_PageLimitStopsEndlessChain(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - stats, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + _, pages, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) - assert.Equal(t, listMaxPages, stats.pages) + assert.Equal(t, listMaxPages, pages) assert.Equal(t, int64(listMaxPages), calls.Load()) } @@ -116,7 +138,7 @@ func TestReconcile_RepeatedTokenEndsWalk(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - _, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + _, _, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) assert.Equal(t, int64(2), calls.Load()) } @@ -124,24 +146,24 @@ func TestReconcile_RepeatedTokenEndsWalk(t *testing.T) { func TestReconcile_GateSkipsSecondReplica(t *testing.T) { srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) store, mr := newTestRedisStore(t) - second, err := NewRedisStore(mr.Addr()) + second, err := contextcache.NewRedisStore(mr.Addr()) require.NoError(t, err) a := resolverWithServer(srv.URL, store) b := resolverWithServer(srv.URL, second) - statsA, err := a.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + statsA, _, err := reconcileOnce(t, a, reconcileAuth) require.NoError(t, err) - assert.False(t, statsA.gated) - statsB, err := b.reconcileOnce(context.Background(), second.(*redisStore), reconcileAuth, time.Minute) + assert.False(t, statsA.Gated) + statsB, _, err := reconcileOnce(t, b, reconcileAuth) require.NoError(t, err) - assert.True(t, statsB.gated) + assert.True(t, statsB.Gated) assert.Equal(t, int64(1), calls.Load()) // Once the interval passes, a replica may reconcile again. mr.FastForward(time.Minute + time.Second) - statsB, err = b.reconcileOnce(context.Background(), second.(*redisStore), reconcileAuth, time.Minute) + statsB, _, err = reconcileOnce(t, b, reconcileAuth) require.NoError(t, err) - assert.False(t, statsB.gated) + assert.False(t, statsB.Gated) } func TestReconcile_GateIsPerProjectRegion(t *testing.T) { @@ -150,11 +172,11 @@ func TestReconcile_GateIsPerProjectRegion(t *testing.T) { r := resolverWithServer(srv.URL, store) other := &fakeGCPAuth{token: "tok", region: "europe-west1", project: "p"} - _, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + _, _, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) - stats, err := r.reconcileOnce(context.Background(), store, other, time.Minute) + stats, _, err := reconcileOnce(t, r, other) require.NoError(t, err) - assert.False(t, stats.gated, "one region's round must not suppress another's") + assert.False(t, stats.Gated, "one region's round must not suppress another's") assert.Equal(t, int64(2), calls.Load()) } @@ -170,11 +192,13 @@ func TestReconcile_SkipsForeignExpiredAndUnparseable(t *testing.T) { store, mr := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - stats, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + stats, _, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) - assert.Equal(t, 4, stats.seen) - assert.Equal(t, 1, stats.written) - assert.Equal(t, 3, stats.skipped) + // 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))) @@ -187,12 +211,12 @@ func TestReconcile_EntryTTLMatchesExpiry(t *testing.T) { store, mr := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - _, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + _, _, err := reconcileOnce(t, r, reconcileAuth) require.NoError(t, err) assert.InDelta(t, (5 * time.Minute).Seconds(), mr.TTL(testKey(1)).Seconds(), 5) } -func TestReconcile_ListFailureReturnsErrorWithoutPanicking(t *testing.T) { +func TestReconcile_ListFailureReturnsError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) @@ -200,42 +224,14 @@ func TestReconcile_ListFailureReturnsErrorWithoutPanicking(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - _, err := r.reconcileOnce(context.Background(), store, reconcileAuth, time.Minute) + _, _, err := reconcileOnce(t, r, reconcileAuth) require.ErrorContains(t, err, "HTTP 500") - // reconcileRound is the caller the loop uses; it must swallow the error. - r.reconcileRound(context.Background(), store, reconcileAuth, time.Minute) } -func TestRunReconciler_NoopStoreReturnsImmediately(t *testing.T) { +func TestNewReconciler_NoopStoreHasNone(t *testing.T) { r := New(nil, nil, nil).(*resolver) - done := make(chan struct{}) - go func() { - r.runReconciler(context.Background(), reconcileAuth, time.Minute) - close(done) - }() - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("runReconciler must return at once for a store that cannot reconcile") - } -} - -func TestRunReconciler_StopsOnContextCancel(t *testing.T) { - srv, _, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) - store, _ := newTestRedisStore(t) - r := resolverWithServer(srv.URL, store) - ctx, cancel := context.WithCancel(context.Background()) - done := make(chan struct{}) - go func() { - r.runReconciler(ctx, reconcileAuth, time.Hour) - close(done) - }() - cancel() - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("runReconciler must return when its context is cancelled") - } + _, ok := r.newReconciler(reconcileAuth, time.Minute) + assert.False(t, ok, "a store that cannot reconcile must not get a reconciler") } func TestIsCacheKey(t *testing.T) { diff --git a/internal/gcpcache/redis_test.go b/internal/gcpcache/redis_test.go index 80ec91932d..64c3ab67d1 100644 --- a/internal/gcpcache/redis_test.go +++ b/internal/gcpcache/redis_test.go @@ -20,87 +20,9 @@ import ( "github.com/stretchr/testify/require" "github.com/envoyproxy/ai-gateway/internal/apischema/openai" + "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 := NewRedisStore(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", entry{cacheName: "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.cacheName) - 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 := entry{cacheName: "projects/p|weird/cachedContents/x", expireTime: expire} - - got, err := decodeEntry(encodeEntry(e)) - require.NoError(t, err) - assert.Equal(t, e.cacheName, got.cacheName) - 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) - } -} - // ----------------------------------------------------------------------- // Cross-replica behavior // ----------------------------------------------------------------------- @@ -114,7 +36,7 @@ func TestResolver_CrossReplica_SecondResolveHitsStore(t *testing.T) { mr := miniredis.RunT(t) newReplica := func() *resolver { - store, err := NewRedisStore(mr.Addr()) + store, err := contextcache.NewRedisStore(mr.Addr()) require.NoError(t, err) return resolverWithServer(fake.srv.URL, store) } @@ -171,7 +93,7 @@ func TestResolver_CrossReplica_ColdPrefixRaceMayDuplicate(t *testing.T) { results := make([]*ResolveResult, replicas) errs := make([]error, replicas) for i := range replicas { - store, err := NewRedisStore(mr.Addr()) + store, err := contextcache.NewRedisStore(mr.Addr()) require.NoError(t, err) r := resolverWithServer(srv.URL, store) done.Add(1) @@ -216,7 +138,7 @@ func TestResolver_StoreUnreachable_FailsOpen(t *testing.T) { fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) mr := miniredis.RunT(t) - store, err := NewRedisStore(mr.Addr()) + store, err := contextcache.NewRedisStore(mr.Addr()) require.NoError(t, err) mr.Close() // Every operation now fails to dial. diff --git a/internal/gcpcache/store.go b/internal/gcpcache/store.go deleted file mode 100644 index bd422e0adf..0000000000 --- a/internal/gcpcache/store.go +++ /dev/null @@ -1,49 +0,0 @@ -// 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" - "time" -) - -// entry is a resolved Google cachedContents entry: the resource name and when it expires. -type entry struct { - cacheName string - expireTime time.Time -} - -// CacheStore is the storage layer behind the resolver, mapping a deterministic cache -// key to the Google cache name it resolved to. -// -// The store is shared across gateway replicas, so a cache name published by one -// replica is reused by all of them. It is a cache, not a lock: replicas that miss the -// same cold prefix at the same moment may each create a Google cache, and the resolver -// tolerates that. Implementations must be safe for concurrent use. -// -// Store errors are never fatal to a request: the resolver treats them as a cache miss -// and proceeds to Google. See the failure policy in the Redis cache store proposal. -type CacheStore 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 -} - -// noopStore is the CacheStore 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. -type noopStore struct{} - -func (noopStore) Get(context.Context, string) (entry, bool, error) { return entry{}, false, nil } - -func (noopStore) Set(context.Context, string, entry, time.Duration) error { return nil } - -var _ CacheStore = noopStore{} From fee488715180335f5f1b525c712321a58d538e8b Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 30 Sep 2026 18:06:55 +0530 Subject: [PATCH 11/17] gcpcache: start and stop the reconciler with the resolver The reconciler existed but nothing ran it. Tie its lifetime to the resolver that owns it. The resolver gains Start and Close behind a new Background interface, so CacheResolver and the request path are unchanged. Start builds the reconciler for the backend's project and region and runs it under a context the resolver owns. It cannot use the LoadConfig context: the config watcher cancels that when each reload finishes, which would stop the reconciler within seconds. Start does nothing if the resolver is already running, has been closed, or has a store that cannot reconcile. Close cancels the context and waits for the goroutine to exit, and is safe to call more than once. LoadConfig calls Start with the backend's GCPAuthHandler. Resolvers carried across a reload are already running and are not restarted. After a reload, LoadConfig closes every resolver that was not carried over, because its backend was removed, lost its cache config, or was repointed at another Redis. Without this, each such reload would leak a goroutine. Server.Close closes all resolvers, and extproc calls it on shutdown after the gRPC server stops. Rounds run every 60s, starting as soon as the reconciler starts. A round against an unreachable Redis fails and is logged at warn level. Signed-off-by: Sivanantham Chinnaiyan --- cmd/extproc/mainlib/main.go | 2 + internal/extproc/server.go | 26 ++++++++++++ internal/extproc/server_test.go | 46 +++++++++++++++++++++ internal/gcpcache/gcpcache.go | 62 ++++++++++++++++++++++++++++ internal/gcpcache/reconciler_test.go | 56 +++++++++++++++++++++++++ 5 files changed, 192 insertions(+) diff --git a/cmd/extproc/mainlib/main.go b/cmd/extproc/mainlib/main.go index c6d23125e6..3cdc12f288 100644 --- a/cmd/extproc/mainlib/main.go +++ b/cmd/extproc/mainlib/main.go @@ -409,6 +409,8 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) { go func() { <-ctx.Done() s.GracefulStop() + // Stop context-cache reconcilers once no request can reach a resolver. + server.Close() shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() diff --git a/internal/extproc/server.go b/internal/extproc/server.go index ebf1318e00..bb2c0e8401 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -127,15 +127,41 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error } resolver = gcpcache.New(nil, store, logger) } + if bg, ok := resolver.(gcpcache.Background); ok { + // Start is a no-op for a resolver already running, so reused resolvers are + // unaffected. It runs under the resolver's own context, not ctx, which is + // cancelled when this reload finishes. + bg.Start(rb.Handler.(filterapi.GCPAuthHandler)) + } live[key] = resolver rb.CacheResolver = resolver } + // Stop the background work of resolvers that were not carried into this config: their + // backend was removed, lost its cache config, or was repointed at another Redis. + for key, resolver := range s.cacheResolvers { + if _, ok := live[key]; ok { + continue + } + if c, ok := resolver.(io.Closer); ok { + _ = c.Close() + } + } s.cacheResolvers = live s.config = newConfig // This is racey, but we don't care. return nil } +// Close stops the background work of every context-cache resolver. It is safe to call +// more than once. +func (s *Server) Close() { + for _, resolver := range s.cacheResolvers { + if c, ok := resolver.(io.Closer); ok { + _ = c.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 a125759721..b144768b81 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -19,6 +19,7 @@ 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" @@ -48,6 +49,7 @@ 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) @@ -57,6 +59,7 @@ func TestServer_LoadConfig(t *testing.T) { // to a GCP backend only when ContextCache is configured with a store URL. func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { s := &Server{} + t.Cleanup(s.Close) config := &filterapi.Config{ Backends: []filterapi.Backend{ @@ -122,6 +125,7 @@ func gcpBackendWithCaching(name, redisURL string) filterapi.Backend { // discard live connections on a hot path that fires whenever any unrelated backend changes. func TestServer_LoadConfig_GCPCacheResolver_ReusedAcrossReloads(t *testing.T) { s := &Server{} + t.Cleanup(s.Close) config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}} require.NoError(t, s.LoadConfig(t.Context(), config)) @@ -136,6 +140,7 @@ func TestServer_LoadConfig_GCPCacheResolver_ReusedAcrossReloads(t *testing.T) { // Repointing a backend at a different Redis must not keep reusing a pool dialing the old one. func TestServer_LoadConfig_GCPCacheResolver_RebuiltWhenRedisURLChanges(t *testing.T) { s := &Server{} + t.Cleanup(s.Close) require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}})) @@ -151,6 +156,7 @@ func TestServer_LoadConfig_GCPCacheResolver_RebuiltWhenRedisURLChanges(t *testin // Resolvers for backends that vanish must not accumulate: the map is rebuilt each reload. func TestServer_LoadConfig_GCPCacheResolver_DroppedWhenBackendDisappears(t *testing.T) { s := &Server{} + t.Cleanup(s.Close) require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ gcpBackendWithCaching("gcp-a", "localhost:6379"), @@ -168,6 +174,7 @@ func TestServer_LoadConfig_GCPCacheResolver_DroppedWhenBackendDisappears(t *test // reload, which would take down backends that have nothing to do with caching. func TestServer_LoadConfig_GCPCacheResolver_BadRedisURLDoesNotFailReload(t *testing.T) { s := &Server{} + t.Cleanup(s.Close) config := &filterapi.Config{Backends: []filterapi.Backend{ gcpBackendWithCaching("gcp-bad-redis", "http://not-a-redis-url"), gcpBackendWithCaching("gcp-ok", "localhost:6379"), @@ -180,6 +187,45 @@ func TestServer_LoadConfig_GCPCacheResolver_BadRedisURLDoesNotFailReload(t *test "an unrelated backend must keep its resolver") } +// A resolver evicted by a reload must have its background reconciler stopped, or every +// reload that drops or repoints a backend would leak a goroutine. +func TestServer_LoadConfig_GCPCacheResolver_EvictedResolverIsClosed(t *testing.T) { + before := goleak.IgnoreCurrent() + s := &Server{} + + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp-a", "localhost:6379"), + gcpBackendWithCaching("gcp-b", "localhost:6379"), + }})) + kept := s.cacheResolvers["gcp-a|localhost:6379"].(io.Closer) + + require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp-a", "localhost:6379"), + }})) + // Closing the one resolver still in the config must leave nothing running. If the + // evicted resolver had not been closed by the reload, its reconciler would remain. + require.NoError(t, kept.Close()) + goleak.VerifyNone(t, before) +} + +// The reload context is cancelled when the reload finishes. Reconcilers must not be tied +// to it, or they would stop seconds after every config load. +func TestServer_LoadConfig_ReconcilerOutlivesReloadContext(t *testing.T) { + before := goleak.IgnoreCurrent() + s := &Server{} + + ctx, cancel := context.WithCancel(t.Context()) + require.NoError(t, s.LoadConfig(ctx, &filterapi.Config{Backends: []filterapi.Backend{ + gcpBackendWithCaching("gcp", "localhost:6379"), + }})) + cancel() + time.Sleep(50 * time.Millisecond) + require.Error(t, goleak.Find(before), "the reconciler 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/gcpcache/gcpcache.go b/internal/gcpcache/gcpcache.go index c396f478f8..10f7166ce5 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/gcpcache/gcpcache.go @@ -32,6 +32,7 @@ import ( "io" "log/slog" "net/http" + "sync" "time" "google.golang.org/genai" @@ -92,6 +93,67 @@ type resolver struct { // logger records store failures, which are otherwise invisible because the request // proceeds normally. logger *slog.Logger + + // mu guards the background reconciler's lifecycle. + mu sync.Mutex + // cancel stops the reconciler; nil until Start runs one. + cancel context.CancelFunc + // done is closed when the reconciler goroutine exits. + done chan struct{} + // closed is set by Close, after which Start does nothing. + closed bool +} + +// reconcileInterval is how often the background reconciler 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 reconcileInterval = 60 * time.Second + +// Background is implemented by resolvers that run background work. It is kept out of +// CacheResolver so that request-path callers need not know about it; LoadConfig reaches +// it through a type assertion. +type Background interface { + // Start begins background reconciliation for gcpAuth's project and region. It runs + // under a context the resolver owns, not the caller's, so it outlives the call. It + // does nothing if already started, if closed, or if the store cannot reconcile. + Start(gcpAuth filterapi.GCPAuthHandler) + // Close stops background work and waits for it to exit. It is idempotent. + io.Closer +} + +var _ Background = (*resolver)(nil) + +// Start implements Background. +func (r *resolver) Start(gcpAuth filterapi.GCPAuthHandler) { + r.mu.Lock() + defer r.mu.Unlock() + if r.closed || r.cancel != nil { + return + } + rc, ok := r.newReconciler(gcpAuth, reconcileInterval) + if !ok { + return + } + ctx, cancel := context.WithCancel(context.Background()) + r.cancel = cancel + r.done = make(chan struct{}) + go func() { + defer close(r.done) + rc.Run(ctx) + }() +} + +// Close implements Background. +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 CacheResolver. diff --git a/internal/gcpcache/reconciler_test.go b/internal/gcpcache/reconciler_test.go index 43cece4889..dc4a2c1a68 100644 --- a/internal/gcpcache/reconciler_test.go +++ b/internal/gcpcache/reconciler_test.go @@ -240,3 +240,59 @@ func TestIsCacheKey(t *testing.T) { assert.False(t, isCacheKey(strings.Repeat("z", 64))) assert.False(t, isCacheKey(testKey(7)[:63])) } + +// ----------------------------------------------------------------------- +// Lifecycle +// ----------------------------------------------------------------------- + +func TestResolver_StartRunsReconcilerAndCloseStopsIt(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(reconcileAuth) + require.Eventually(t, func() bool { return mr.Exists(testKey(1)) }, 2*time.Second, 10*time.Millisecond, + "Start must run a reconcile round") + + require.NoError(t, r.Close()) + select { + case <-r.done: + default: + t.Fatal("Close must wait for the reconciler 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(reconcileAuth) + first := r.done + r.Start(reconcileAuth) + assert.Equal(t, first, r.done, "a second Start must not launch another reconciler") + 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(reconcileAuth) + assert.Nil(t, r.cancel, "Start after Close must do nothing") + assert.Equal(t, int64(0), calls.Load()) +} + +func TestResolver_StartWithNoopStoreDoesNothing(t *testing.T) { + r := New(nil, nil, nil).(*resolver) + r.Start(reconcileAuth) + assert.Nil(t, r.cancel) + require.NoError(t, r.Close()) +} From bf2d1a429cba0502309e375c9cc010571de2b0be Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Sun, 4 Oct 2026 18:21:51 +0530 Subject: [PATCH 12/17] contextcache: restructure packages, refresh credentials, harden cache path Package layout: - Move gcpcache to internal/contextcache/gcpcache and the Redis store to internal/contextcache/redis (redis.NewStore). contextcache itself now holds only interfaces and shared types, so filterapi can import it without linking go-redis. - Move CacheResolver and ResolveResult into contextcache. Resolve no longer takes an auth handler; the GCP resolver uses the credentials passed to New. - Move CacheSyncer into contextcache with Start() taking no arguments, and fold the separate reconcile-store interface into Store. NoopStore implements the sync methods as no-ops. - Rename the reconciler to syncer and its gate prefix to gcpcache:sync:. - Type filterapi.RuntimeBackend.CacheResolver as contextcache.CacheResolver instead of any. Behavior: - Remove the sync walk's fixed page cap; it now follows nextPageToken until it is empty. - Refresh credentials on reused resolvers: LoadConfig calls SetAuth on every reload because the controller rotates the access token. Resolve and the syncer read the current credentials under a lock. - Include GCP project and region in the cache key, so backends in different locations sharing one Redis do not reuse each other's caches. Existing keys change, so each prefix misses once after deploy. - Skip resolution when a request names its own cachedContent, so no billed cache is created for a request the translator rejects. - Only bill a cache write when the cache was created. - Send expireTime on create only when set: it is now a pointer, since omitempty does not omit a zero time.Time, which was being sent alongside ttl. Naming: - Rename translator GCPCacheSetter/GCPCacheResult to ContextCacheSetter/ContextCacheResult, and pendingCacheResult to cacheResult. Docs and tests: - Drop the claim that the shared store prevents duplicate creates from the ContextCacheSpec.URL and filterapi godoc, and regenerate the CRD and API docs. - Fix the crdcel fixtures, which still used the removed gcpContextCaching field and so did not exercise context caching, and add a missing-url case. - Fix a resolver test that passed only because it had no credentials. Signed-off-by: Sivanantham Chinnaiyan --- api/v1alpha1/ai_service_backend.go | 3 +- api/v1beta1/ai_service_backend.go | 3 +- cmd/extproc/mainlib/main.go | 2 +- internal/apischema/gcp/gcp.go | 4 +- .../{ => contextcache}/gcpcache/gcpcache.go | 169 ++++++++---------- .../gcpcache/gcpcache_test.go | 104 +++++++---- .../{ => contextcache}/gcpcache/redis_test.go | 22 ++- .../gcpcache/syncer.go} | 69 +++---- .../gcpcache/syncer_test.go} | 159 ++++++++-------- .../contextcache/{redis.go => redis/store.go} | 52 +++--- .../{redis_test.go => redis/store_test.go} | 10 +- internal/contextcache/resolver.go | 43 +++++ internal/contextcache/store.go | 31 ++-- .../contextcache/{reconciler.go => syncer.go} | 44 +++-- .../{reconciler_test.go => syncer_test.go} | 94 ++++++++-- internal/extproc/processor_impl.go | 46 ++--- internal/extproc/processor_impl_test.go | 135 +++++++------- internal/extproc/server.go | 56 +++--- internal/extproc/server_test.go | 29 ++- internal/filterapi/filterconfig.go | 8 +- internal/filterapi/runtime.go | 9 +- internal/translator/openai_gcpvertexai.go | 38 ++-- .../translator/openai_gcpvertexai_test.go | 51 +++++- internal/translator/translator.go | 25 +-- ...teway.envoyproxy.io_aiservicebackends.yaml | 6 +- site/docs/api/api.mdx | 4 +- tests/crdcel/main_test.go | 12 +- ...nabled.yaml => context-cache-enabled.yaml} | 4 +- ...tl.yaml => context-cache-invalid-ttl.yaml} | 4 +- .../context-cache-missing-url.yaml | 20 +++ 30 files changed, 768 insertions(+), 488 deletions(-) rename internal/{ => contextcache}/gcpcache/gcpcache.go (75%) rename internal/{ => contextcache}/gcpcache/gcpcache_test.go (86%) rename internal/{ => contextcache}/gcpcache/redis_test.go (90%) rename internal/{gcpcache/reconciler.go => contextcache/gcpcache/syncer.go} (72%) rename internal/{gcpcache/reconciler_test.go => contextcache/gcpcache/syncer_test.go} (67%) rename internal/contextcache/{redis.go => redis/store.go} (58%) rename internal/contextcache/{redis_test.go => redis/store_test.go} (89%) create mode 100644 internal/contextcache/resolver.go rename internal/contextcache/{reconciler.go => syncer.go} (74%) rename internal/contextcache/{reconciler_test.go => syncer_test.go} (51%) rename tests/crdcel/testdata/aiservicebackends/{gcp-context-caching-enabled.yaml => context-cache-enabled.yaml} (87%) rename tests/crdcel/testdata/aiservicebackends/{gcp-context-caching-invalid-ttl.yaml => context-cache-invalid-ttl.yaml} (88%) create mode 100644 tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml diff --git a/api/v1alpha1/ai_service_backend.go b/api/v1alpha1/ai_service_backend.go index 0178f59d55..c586e99a81 100644 --- a/api/v1alpha1/ai_service_backend.go +++ b/api/v1alpha1/ai_service_backend.go @@ -103,8 +103,7 @@ type ContextCacheSpec struct { DefaultTTL string `json:"defaultTTL,omitempty"` // URL locates the shared Redis store that records which cachedContents entry a given - // request prefix resolved to. Sharing that state across gateway replicas is what - // prevents two replicas from independently creating a cache for the same prefix. + // request prefix resolved to. // // It is given either as a bare "host:port" // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. diff --git a/api/v1beta1/ai_service_backend.go b/api/v1beta1/ai_service_backend.go index 46c20b3501..f45ff7bb52 100644 --- a/api/v1beta1/ai_service_backend.go +++ b/api/v1beta1/ai_service_backend.go @@ -118,8 +118,7 @@ type ContextCacheSpec struct { DefaultTTL string `json:"defaultTTL,omitempty"` // URL locates the shared Redis store that records which cachedContents entry a given - // request prefix resolved to. Sharing that state across gateway replicas is what - // prevents two replicas from independently creating a cache for the same prefix. + // request prefix resolved to. // // It is given either as a bare "host:port" // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. diff --git a/cmd/extproc/mainlib/main.go b/cmd/extproc/mainlib/main.go index 156481b69b..317efb48d1 100644 --- a/cmd/extproc/mainlib/main.go +++ b/cmd/extproc/mainlib/main.go @@ -425,7 +425,7 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) { go func() { <-ctx.Done() s.GracefulStop() - // Stop context-cache reconcilers once no request can reach a resolver. + // Stop context-cache syncers once no request can reach a resolver. server.Close() shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) diff --git a/internal/apischema/gcp/gcp.go b/internal/apischema/gcp/gcp.go index 3f5c73027e..b8b2a62743 100644 --- a/internal/apischema/gcp/gcp.go +++ b/internal/apischema/gcp/gcp.go @@ -189,7 +189,9 @@ type CreateCachedContent struct { // 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. - ExpireTime time.Time `json:"expireTime,omitempty"` + // 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. diff --git a/internal/gcpcache/gcpcache.go b/internal/contextcache/gcpcache/gcpcache.go similarity index 75% rename from internal/gcpcache/gcpcache.go rename to internal/contextcache/gcpcache/gcpcache.go index 10f7166ce5..a4abcd8372 100644 --- a/internal/gcpcache/gcpcache.go +++ b/internal/contextcache/gcpcache/gcpcache.go @@ -13,14 +13,6 @@ // 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. -// -// 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. -// -// The CacheResolver interface allows the implementation to be replaced with an external -// cache service in the future without changing the request-path callers. package gcpcache import ( @@ -53,35 +45,8 @@ const ( gcpCachedContentsBasePath = "https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/cachedContents" ) -// ResolveResult holds the outcome of a successful cache resolution. -type ResolveResult struct { - // CacheName is the full Google resource name of the resolved or created cache entry. - // Format: "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 Google's create response). - // Only populated when Created is true. - TokenCount int - // ExpireTime is the cache expiration time reported by Google. - ExpireTime time.Time -} - -// CacheResolver resolves or creates a GCP Vertex AI cached content entry for requests -// that carry Anthropic-style cache_control markers. -type CacheResolver interface { - // Resolve inspects openAIReq for cache_control markers, resolves or creates the - // corresponding Google cachedContents entry, and returns the result. - // The caller is responsible for injecting ResolveResult.CacheName as the - // cachedContent field on the Gemini request and replacing the request messages - // with ResolveResult.Messages. - Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) -} - -// resolver is the default CacheResolver implementation. -type resolver struct { +// 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 @@ -94,46 +59,40 @@ type resolver struct { // proceeds normally. logger *slog.Logger - // mu guards the background reconciler's lifecycle. + // mu guards auth and the background syncer's lifecycle. mu sync.Mutex - // cancel stops the reconciler; nil until Start runs one. + // 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 reconciler goroutine exits. + // done is closed when the syncer goroutine exits. done chan struct{} // closed is set by Close, after which Start does nothing. closed bool } -// reconcileInterval is how often the background reconciler runs a round. Between rounds, +// 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 reconcileInterval = 60 * time.Second - -// Background is implemented by resolvers that run background work. It is kept out of -// CacheResolver so that request-path callers need not know about it; LoadConfig reaches -// it through a type assertion. -type Background interface { - // Start begins background reconciliation for gcpAuth's project and region. It runs - // under a context the resolver owns, not the caller's, so it outlives the call. It - // does nothing if already started, if closed, or if the store cannot reconcile. - Start(gcpAuth filterapi.GCPAuthHandler) - // Close stops background work and waits for it to exit. It is idempotent. - io.Closer -} +const syncInterval = 60 * time.Second + +var _ contextcache.CacheResolver = (*Resolver)(nil) -var _ Background = (*resolver)(nil) +var _ contextcache.CacheSyncer = (*Resolver)(nil) -// Start implements Background. -func (r *resolver) Start(gcpAuth filterapi.GCPAuthHandler) { +// 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 { - return - } - rc, ok := r.newReconciler(gcpAuth, reconcileInterval) - if !ok { + 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{}) @@ -143,8 +102,25 @@ func (r *resolver) Start(gcpAuth filterapi.GCPAuthHandler) { }() } -// Close implements Background. -func (r *resolver) Close() error { +// 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 @@ -156,12 +132,13 @@ func (r *resolver) Close() error { return nil } -// New creates a new CacheResolver. +// New creates a new Resolver. // -// httpClient is used for calls to the Google cachedContents API; pass nil for a default. +// 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. logger may be nil. -func New(httpClient *http.Client, store contextcache.Store, logger *slog.Logger) CacheResolver { +// 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, logger *slog.Logger) *Resolver { if httpClient == nil { httpClient = &http.Client{Timeout: 30 * time.Second} } @@ -171,11 +148,16 @@ func New(httpClient *http.Client, store contextcache.Store, logger *slog.Logger) if logger == nil { logger = slog.New(slog.DiscardHandler) } - return &resolver{httpClient: httpClient, store: store, logger: logger} + return &Resolver{httpClient: httpClient, store: store, logger: logger, auth: auth} } -// Resolve implements CacheResolver. -func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler) (*ResolveResult, error) { +// 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 { @@ -201,16 +183,16 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion } // Compute a deterministic cache key. - cacheKey, err := computeCacheKey(openAIReq.Model, contents, systemInstruction, geminiTools) + 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. Google-side failures below are a different matter and do fail fast. + // request. GCP-side failures below are a different matter and do fail fast. if e, ok := r.storeGet(ctx, cacheKey); ok { - return &ResolveResult{ + return &contextcache.ResolveResult{ CacheName: e.Name, Messages: remainderMessages, ExpireTime: e.ExpireTime, @@ -219,7 +201,7 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // Store miss — create against GCP. See resolveUncached for how concurrent // resolutions of the same key behave. - result, err := r.resolveUncached(ctx, openAIReq, gcpAuth, cacheKey, ttl, contents, systemInstruction, geminiTools) + result, err := r.resolveUncached(ctx, openAIReq, auth, cacheKey, ttl, contents, systemInstruction, geminiTools) if err != nil { return nil, err } @@ -230,7 +212,7 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion return result, nil } -// resolveUncached creates a Google cachedContents entry for a cache key that missed the +// 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. // @@ -238,7 +220,7 @@ func (r *resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // 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( +func (r *Resolver) resolveUncached( ctx context.Context, openAIReq *openai.ChatCompletionRequest, gcpAuth filterapi.GCPAuthHandler, @@ -246,7 +228,7 @@ func (r *resolver) resolveUncached( contents []genai.Content, systemInstruction *genai.Content, geminiTools []genai.Tool, -) (*ResolveResult, error) { +) (*contextcache.ResolveResult, error) { region := gcpAuth.GCPRegion() project := gcpAuth.GCPProject() @@ -259,7 +241,7 @@ func (r *resolver) resolveUncached( // 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 Google caches. + // 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 { @@ -267,7 +249,7 @@ func (r *resolver) resolveUncached( } r.storeSet(ctx, cacheKey, contextcache.Entry{Name: created, ExpireTime: expireTime}) - return &ResolveResult{ + return &contextcache.ResolveResult{ CacheName: created, Created: true, TokenCount: tokenCount, @@ -280,15 +262,15 @@ func (r *resolver) resolveUncached( // // 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 Google as the source of -// truth. Google-side failures are propagated, because those are actionable by the +// 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) { +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", @@ -306,9 +288,9 @@ func (r *resolver) storeGet(ctx context.Context, key string) (contextcache.Entry return e, true } -// storeSet records a resolved cache name, expiring it with the Google entry itself so +// 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) { +func (r *Resolver) storeSet(ctx context.Context, key string, e contextcache.Entry) { ttl := time.Until(e.ExpireTime) if ttl <= 0 { return @@ -446,6 +428,11 @@ func computeKeyInputs(model string, cachedMessages []openai.ChatCompletionMessag // ----------------------------------------------------------------------- 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"` @@ -453,10 +440,12 @@ type cacheKeyInput struct { } // computeCacheKey generates a deterministic SHA-256 hex digest over the -// (model, contents, systemInstruction, tools) tuple. The digest is used as the -// Google cachedContents displayName. -func computeCacheKey(model string, contents []genai.Content, systemInstruction *genai.Content, tools []genai.Tool) (string, error) { +// (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, @@ -471,12 +460,12 @@ func computeCacheKey(model string, contents []genai.Content, systemInstruction * } // ----------------------------------------------------------------------- -// Google cachedContents REST API calls +// 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( +func (r *Resolver) createCache( ctx context.Context, baseURL, accessToken, model, region, project, cacheKey string, contents []genai.Content, diff --git a/internal/gcpcache/gcpcache_test.go b/internal/contextcache/gcpcache/gcpcache_test.go similarity index 86% rename from internal/gcpcache/gcpcache_test.go rename to internal/contextcache/gcpcache/gcpcache_test.go index f453f1604e..6500eeaeb0 100644 --- a/internal/gcpcache/gcpcache_test.go +++ b/internal/contextcache/gcpcache/gcpcache_test.go @@ -147,7 +147,7 @@ func (t *redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) // 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 { +func resolverWithServer(srvURL string, store ...contextcache.Store) *Resolver { host := srvURL[len("http://"):] var s contextcache.Store if len(store) > 0 { @@ -156,7 +156,7 @@ func resolverWithServer(srvURL string, store ...contextcache.Store) *resolver { return New(&http.Client{ Transport: &redirectTransport{fakeHost: host}, Timeout: 5 * time.Second, - }, s, nil).(*resolver) + }, s, syncAuth, nil) } // memStore is an in-process contextcache.Store for tests. @@ -195,6 +195,21 @@ func (m *memStore) Set(_ context.Context, key string, e contextcache.Entry, _ ti 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() @@ -208,7 +223,7 @@ func computeKeyForRequest(t *testing.T, req *openai.ChatCompletionRequest) strin require.GreaterOrEqual(t, bp, 0) contents, sys, err := computeKeyInputs(req.Model, req.Messages[:bp+1]) require.NoError(t, err) - key, err := computeCacheKey(req.Model, contents, sys, nil) + key, err := computeCacheKey(syncAuth.project, syncAuth.region, req.Model, contents, sys, nil) require.NoError(t, err) return key } @@ -327,9 +342,9 @@ func TestComputeCacheKey_Deterministic(t *testing.T) { contents, sys, err := computeKeyInputs("gemini-1.5-pro", msgs) require.NoError(t, err) - k1, err := computeCacheKey("gemini-1.5-pro", contents, sys, nil) + k1, err := computeCacheKey("p", "r", "gemini-1.5-pro", contents, sys, nil) require.NoError(t, err) - k2, err := computeCacheKey("gemini-1.5-pro", contents, sys, nil) + 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") @@ -339,25 +354,58 @@ 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("gemini-1.5-pro", c, s, nil) - k2, _ := computeCacheKey("gemini-2.0-flash", c, s, nil) + 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) { - r := New(nil, nil, nil).(*resolver) + // 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")}, } - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) { @@ -366,7 +414,6 @@ func TestResolver_CacheMiss_Creates(t *testing.T) { fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) r := resolverWithServer(fake.srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -375,7 +422,7 @@ func TestResolver_CacheMiss_Creates(t *testing.T) { }, } - res, err := r.Resolve(context.Background(), req, auth) + 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) @@ -390,7 +437,7 @@ func TestResolver_CacheMiss_Creates(t *testing.T) { // 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 reconciler's job. +// the syncer's job. func TestResolver_RequestPath_NeverLists(t *testing.T) { req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", @@ -417,9 +464,8 @@ func TestResolver_RequestPath_NeverLists(t *testing.T) { createResp := `{"name":"projects/p/locations/us-central1/cachedContents/new","expireTime":"` + expireISO + `"}` fake := newFakeCacheServer(t, string(listBody), createResp, http.StatusOK) r := resolverWithServer(fake.srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) @@ -443,8 +489,7 @@ func TestResolver_StoreHit_SkipsGoogleAPICalls(t *testing.T) { key := computeKeyForRequest(t, req) store.seed(key, "projects/p/locations/r/cachedContents/store-hit", time.Now().Add(10*time.Minute)) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) @@ -471,8 +516,7 @@ func TestResolver_StoreEntryExpiring_Refetches(t *testing.T) { // 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)) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) @@ -493,7 +537,6 @@ func TestResolver_CreateExpiryMissing_Fails(t *testing.T) { `{"name":"projects/p/locations/us-central1/cachedContents/x"}`, http.StatusOK) store := newMemStore() r := resolverWithServer(fake.srv.URL, store) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -502,7 +545,7 @@ func TestResolver_CreateExpiryMissing_Fails(t *testing.T) { }, } - res, err := r.Resolve(context.Background(), req, auth) + res, err := r.Resolve(context.Background(), req) require.Error(t, err) assert.Nil(t, res) assert.Contains(t, err.Error(), "no expireTime") @@ -513,7 +556,6 @@ 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) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -522,7 +564,7 @@ func TestResolver_CreateExpiryUnparseable_Fails(t *testing.T) { }, } - res, err := r.Resolve(context.Background(), req, auth) + res, err := r.Resolve(context.Background(), req) require.Error(t, err) assert.Nil(t, res) } @@ -534,7 +576,6 @@ func TestResolver_CreateFailure_ReturnsError(t *testing.T) { http.StatusUnprocessableEntity, ) r := resolverWithServer(fake.srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -542,7 +583,7 @@ func TestResolver_CreateFailure_ReturnsError(t *testing.T) { userMsg("q"), }, } - _, err := r.Resolve(context.Background(), req, auth) + _, err := r.Resolve(context.Background(), req) require.Error(t, err) assert.Contains(t, err.Error(), "422") } @@ -567,7 +608,6 @@ func TestResolver_TTLDefault_SentToGoogle(t *testing.T) { defer srv.Close() r := resolverWithServer(srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -575,12 +615,16 @@ func TestResolver_TTLDefault_SentToGoogle(t *testing.T) { userMsg("q"), }, } - _, err := r.Resolve(context.Background(), req, auth) + _, err := r.Resolve(context.Background(), req) require.NoError(t, err) var body gcp.CreateCachedContent require.NoError(t, json.Unmarshal(capturedBody, &body)) assert.Equal(t, defaultTTL, 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) { @@ -603,7 +647,6 @@ func TestResolver_TTLOverride_SentToGoogle(t *testing.T) { defer srv.Close() r := resolverWithServer(srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} req := &openai.ChatCompletionRequest{ Model: "gemini-1.5-pro", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -611,7 +654,7 @@ func TestResolver_TTLOverride_SentToGoogle(t *testing.T) { userMsg("q"), }, } - _, err := r.Resolve(context.Background(), req, auth) + _, err := r.Resolve(context.Background(), req) require.NoError(t, err) var body gcp.CreateCachedContent @@ -660,8 +703,7 @@ func TestResolver_DuplicateCreateRace_KeepsOwnCreate(t *testing.T) { defer srv.Close() r := resolverWithServer(srv.URL) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) diff --git a/internal/gcpcache/redis_test.go b/internal/contextcache/gcpcache/redis_test.go similarity index 90% rename from internal/gcpcache/redis_test.go rename to internal/contextcache/gcpcache/redis_test.go index 64c3ab67d1..7630660994 100644 --- a/internal/gcpcache/redis_test.go +++ b/internal/contextcache/gcpcache/redis_test.go @@ -21,6 +21,7 @@ import ( "github.com/envoyproxy/ai-gateway/internal/apischema/openai" "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/contextcache/redis" ) // ----------------------------------------------------------------------- @@ -35,21 +36,20 @@ func TestResolver_CrossReplica_SecondResolveHitsStore(t *testing.T) { fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) mr := miniredis.RunT(t) - newReplica := func() *resolver { - store, err := contextcache.NewRedisStore(mr.Addr()) + newReplica := func() *Resolver { + store, err := redis.NewStore(mr.Addr()) require.NoError(t, err) return resolverWithServer(fake.srv.URL, store) } a, b := newReplica(), newReplica() - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - resA, err := a.Resolve(context.Background(), crossReplicaRequest(), auth) + 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(), auth) + resB, err := b.Resolve(context.Background(), crossReplicaRequest()) require.NoError(t, err) require.NotNil(t, resB) assert.Equal(t, resA.CacheName, resB.CacheName) @@ -87,19 +87,18 @@ func TestResolver_CrossReplica_ColdPrefixRaceMayDuplicate(t *testing.T) { t.Cleanup(srv.Close) mr := miniredis.RunT(t) - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} var done sync.WaitGroup - results := make([]*ResolveResult, replicas) + results := make([]*contextcache.ResolveResult, replicas) errs := make([]error, replicas) for i := range replicas { - store, err := contextcache.NewRedisStore(mr.Addr()) + 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(), auth) + results[i], errs[i] = r.Resolve(context.Background(), crossReplicaRequest()) }() } for range replicas { @@ -138,7 +137,7 @@ func TestResolver_StoreUnreachable_FailsOpen(t *testing.T) { fake := newFakeCacheServer(t, `{"cachedContents":[]}`, createResp, http.StatusOK) mr := miniredis.RunT(t) - store, err := contextcache.NewRedisStore(mr.Addr()) + store, err := redis.NewStore(mr.Addr()) require.NoError(t, err) mr.Close() // Every operation now fails to dial. @@ -150,9 +149,8 @@ func TestResolver_StoreUnreachable_FailsOpen(t *testing.T) { userMsg("Hello"), }, } - auth := &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} - res, err := r.Resolve(context.Background(), req, auth) + 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) diff --git a/internal/gcpcache/reconciler.go b/internal/contextcache/gcpcache/syncer.go similarity index 72% rename from internal/gcpcache/reconciler.go rename to internal/contextcache/gcpcache/syncer.go index b55f569b61..3a725f2243 100644 --- a/internal/gcpcache/reconciler.go +++ b/internal/contextcache/gcpcache/syncer.go @@ -21,26 +21,21 @@ import ( "github.com/envoyproxy/ai-gateway/internal/json" ) -// This file supplies the GCP side of reconciliation: a contextcache.Source that walks +// 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.Reconciler. +// 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 - // listMaxPages bounds one round's walk, so a misbehaving nextPageToken chain cannot - // loop forever. Entries past the bound are picked up by the request path creating. - listMaxPages = 10 - - // reconcileGatePrefix namespaces gate keys. Cache keys are 64 hex characters, so the - // prefix cannot collide with them. It is unchanged from before the store moved to - // contextcache, so old and new replicas share gates during a rolling deploy. - reconcileGatePrefix = "gcpcache:reconcile:" + // 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 reconciler uses. +// cachedContentItem is the subset of a cachedContents list item the syncer uses. type cachedContentItem struct { Name string `json:"name"` DisplayName string `json:"displayName"` @@ -52,41 +47,54 @@ type listResponse struct { NextPageToken string `json:"nextPageToken"` } -// newReconciler returns a reconciler for gcpAuth's project and region. It reports false -// when the store cannot back one, such as the no-op store. -func (r *resolver) newReconciler(gcpAuth filterapi.GCPAuthHandler, interval time.Duration) (*contextcache.Reconciler, bool) { - rs, ok := r.store.(contextcache.ReconcileStore) - if !ok { - return nil, false - } - return &contextcache.Reconciler{ - Store: rs, +// 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, - }, true + } } // listSource is a contextcache.Source over cachedContents.list. type listSource struct { - r *resolver + 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 { - return reconcileGatePrefix + s.auth.GCPProject() + "/" + s.auth.GCPRegion() + 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 - region, project := s.auth.GCPRegion(), s.auth.GCPProject() + auth := s.creds() + if auth == nil { + return nil + } + region, project := auth.GCPRegion(), auth.GCPProject() - token, err := s.auth.GCPTokenSource().Token() + token, err := auth.GCPTokenSource().Token() if err != nil { return fmt.Errorf("get GCP access token: %w", err) } @@ -96,7 +104,7 @@ func (s *listSource) List(ctx context.Context, add func(string, contextcache.Ent } pageToken := "" - for s.pages < listMaxPages { + for { lr, err := s.r.listPage(ctx, base, token.AccessToken, pageToken) if err != nil { return err @@ -107,16 +115,11 @@ func (s *listSource) List(ctx context.Context, add func(string, contextcache.Ent add(key, e) } } - // An empty token ends the list. A repeated one would loop, so treat it as the end. - if lr.NextPageToken == "" || lr.NextPageToken == pageToken { + if lr.NextPageToken == "" { return nil } pageToken = lr.NextPageToken } - s.r.logger.Warn("gcpcache: reconcile stopped at page limit; later entries were not read", - slog.String("project", project), slog.String("region", region), - slog.Int("maxPages", listMaxPages)) - return nil } // toEntry converts a list item. It reports false for items the gateway did not create @@ -136,7 +139,7 @@ func (s *listSource) toEntry(item cachedContentItem) (string, contextcache.Entry // 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) { +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)) diff --git a/internal/gcpcache/reconciler_test.go b/internal/contextcache/gcpcache/syncer_test.go similarity index 67% rename from internal/gcpcache/reconciler_test.go rename to internal/contextcache/gcpcache/syncer_test.go index dc4a2c1a68..ac0fe22877 100644 --- a/internal/gcpcache/reconciler_test.go +++ b/internal/contextcache/gcpcache/syncer_test.go @@ -20,25 +20,25 @@ import ( "github.com/stretchr/testify/require" "github.com/envoyproxy/ai-gateway/internal/contextcache" + "github.com/envoyproxy/ai-gateway/internal/contextcache/redis" ) -var reconcileAuth = &fakeGCPAuth{token: "tok", region: "us-central1", project: "p"} +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.ReconcileStore, *miniredis.Miniredis) { +func newTestRedisStore(t *testing.T) (contextcache.Store, *miniredis.Miniredis) { t.Helper() mr := miniredis.RunT(t) - s, err := contextcache.NewRedisStore(mr.Addr()) + s, err := redis.NewStore(mr.Addr()) require.NoError(t, err) return s, mr } -// reconcileOnce runs one round for auth against r's store and returns its stats and the +// syncOnce runs one round for auth against r's store and returns its stats and the // number of list pages read. -func reconcileOnce(t *testing.T, r *resolver, auth *fakeGCPAuth) (contextcache.Stats, int, error) { +func syncOnce(t *testing.T, r *Resolver, auth *fakeGCPAuth) (contextcache.Stats, int, error) { t.Helper() - rc, ok := r.newReconciler(auth, time.Minute) - require.True(t, ok) + rc := r.newSyncer(auth, time.Minute) stats, err := rc.RunOnce(context.Background()) return stats, rc.Source.(*listSource).pages, err } @@ -57,7 +57,7 @@ func newListServer(t *testing.T, pages map[string]string) (*httptest.Server, *at var queries []string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { - t.Errorf("reconciler issued %s; it must only list", r.Method) + t.Errorf("syncer issued %s; it must only list", r.Method) } calls.Add(1) queries = append(queries, r.URL.RawQuery) @@ -73,7 +73,7 @@ func newListServer(t *testing.T, pages map[string]string) (*httptest.Server, *at return srv, &calls, &queries } -func TestReconcile_MultiPageWalkWritesEveryEntry(t *testing.T) { +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"}`, @@ -82,7 +82,7 @@ func TestReconcile_MultiPageWalkWritesEveryEntry(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - stats, pages, err := reconcileOnce(t, r, reconcileAuth) + stats, pages, err := syncOnce(t, r, syncAuth) require.NoError(t, err) assert.Equal(t, 2, pages) assert.Equal(t, 2, stats.Written) @@ -98,7 +98,7 @@ func TestReconcile_MultiPageWalkWritesEveryEntry(t *testing.T) { } } -func TestReconcile_DoesNotOverwriteExistingEntry(t *testing.T) { +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) + `]}`, @@ -107,80 +107,51 @@ func TestReconcile_DoesNotOverwriteExistingEntry(t *testing.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 := reconcileOnce(t, r, reconcileAuth) + 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, "reconciler must not clobber a request's write") + assert.Equal(t, "c/published", e.Name, "syncer must not clobber a request's write") } -func TestReconcile_PageLimitStopsEndlessChain(t *testing.T) { - var calls atomic.Int64 - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - n := calls.Add(1) - _, _ = fmt.Fprintf(w, `{"cachedContents":[],"nextPageToken":"t%d"}`, n) - })) - t.Cleanup(srv.Close) - store, _ := newTestRedisStore(t) - r := resolverWithServer(srv.URL, store) - - _, pages, err := reconcileOnce(t, r, reconcileAuth) - require.NoError(t, err) - assert.Equal(t, listMaxPages, pages) - assert.Equal(t, int64(listMaxPages), calls.Load()) -} - -func TestReconcile_RepeatedTokenEndsWalk(t *testing.T) { - srv, calls, _ := newListServer(t, map[string]string{ - "": `{"cachedContents":[],"nextPageToken":"same"}`, - "same": `{"cachedContents":[],"nextPageToken":"same"}`, - }) - store, _ := newTestRedisStore(t) - r := resolverWithServer(srv.URL, store) - - _, _, err := reconcileOnce(t, r, reconcileAuth) - require.NoError(t, err) - assert.Equal(t, int64(2), calls.Load()) -} - -func TestReconcile_GateSkipsSecondReplica(t *testing.T) { +func TestSync_GateSkipsSecondReplica(t *testing.T) { srv, calls, _ := newListServer(t, map[string]string{"": `{"cachedContents":[]}`}) store, mr := newTestRedisStore(t) - second, err := contextcache.NewRedisStore(mr.Addr()) + second, err := redis.NewStore(mr.Addr()) require.NoError(t, err) a := resolverWithServer(srv.URL, store) b := resolverWithServer(srv.URL, second) - statsA, _, err := reconcileOnce(t, a, reconcileAuth) + statsA, _, err := syncOnce(t, a, syncAuth) require.NoError(t, err) assert.False(t, statsA.Gated) - statsB, _, err := reconcileOnce(t, b, reconcileAuth) + 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 reconcile again. + // Once the interval passes, a replica may sync again. mr.FastForward(time.Minute + time.Second) - statsB, _, err = reconcileOnce(t, b, reconcileAuth) + statsB, _, err = syncOnce(t, b, syncAuth) require.NoError(t, err) assert.False(t, statsB.Gated) } -func TestReconcile_GateIsPerProjectRegion(t *testing.T) { +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 := reconcileOnce(t, r, reconcileAuth) + _, _, err := syncOnce(t, r, syncAuth) require.NoError(t, err) - stats, _, err := reconcileOnce(t, r, other) + 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 TestReconcile_SkipsForeignExpiredAndUnparseable(t *testing.T) { +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 @@ -192,7 +163,7 @@ func TestReconcile_SkipsForeignExpiredAndUnparseable(t *testing.T) { store, mr := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - stats, _, err := reconcileOnce(t, r, reconcileAuth) + 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. @@ -205,18 +176,18 @@ func TestReconcile_SkipsForeignExpiredAndUnparseable(t *testing.T) { assert.True(t, mr.Exists(testKey(3))) } -func TestReconcile_EntryTTLMatchesExpiry(t *testing.T) { +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 := reconcileOnce(t, r, reconcileAuth) + _, _, err := syncOnce(t, r, syncAuth) require.NoError(t, err) assert.InDelta(t, (5 * time.Minute).Seconds(), mr.TTL(testKey(1)).Seconds(), 5) } -func TestReconcile_ListFailureReturnsError(t *testing.T) { +func TestSync_ListFailureReturnsError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) @@ -224,14 +195,19 @@ func TestReconcile_ListFailureReturnsError(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - _, _, err := reconcileOnce(t, r, reconcileAuth) + _, _, err := syncOnce(t, r, syncAuth) require.ErrorContains(t, err, "HTTP 500") } -func TestNewReconciler_NoopStoreHasNone(t *testing.T) { - r := New(nil, nil, nil).(*resolver) - _, ok := r.newReconciler(reconcileAuth, time.Minute) - assert.False(t, ok, "a store that cannot reconcile must not get a reconciler") +// 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) { @@ -245,7 +221,7 @@ func TestIsCacheKey(t *testing.T) { // Lifecycle // ----------------------------------------------------------------------- -func TestResolver_StartRunsReconcilerAndCloseStopsIt(t *testing.T) { +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) + `]}`, @@ -253,15 +229,15 @@ func TestResolver_StartRunsReconcilerAndCloseStopsIt(t *testing.T) { store, mr := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - r.Start(reconcileAuth) + r.Start() require.Eventually(t, func() bool { return mr.Exists(testKey(1)) }, 2*time.Second, 10*time.Millisecond, - "Start must run a reconcile round") + "Start must run a sync round") require.NoError(t, r.Close()) select { case <-r.done: default: - t.Fatal("Close must wait for the reconciler goroutine to exit") + t.Fatal("Close must wait for the syncer goroutine to exit") } assert.Equal(t, int64(1), calls.Load()) } @@ -271,10 +247,10 @@ func TestResolver_StartIsIdempotent(t *testing.T) { store, _ := newTestRedisStore(t) r := resolverWithServer(srv.URL, store) - r.Start(reconcileAuth) + r.Start() first := r.done - r.Start(reconcileAuth) - assert.Equal(t, first, r.done, "a second Start must not launch another reconciler") + r.Start() + assert.Equal(t, first, r.done, "a second Start must not launch another syncer") require.NoError(t, r.Close()) } @@ -285,14 +261,51 @@ func TestResolver_CloseIsIdempotentAndBlocksLaterStart(t *testing.T) { require.NoError(t, r.Close(), "Close before Start must be safe") require.NoError(t, r.Close()) - r.Start(reconcileAuth) + r.Start() assert.Nil(t, r.cancel, "Start after Close must do nothing") assert.Equal(t, int64(0), calls.Load()) } -func TestResolver_StartWithNoopStoreDoesNothing(t *testing.T) { - r := New(nil, nil, nil).(*resolver) - r.Start(reconcileAuth) - assert.Nil(t, r.cancel) +func TestResolver_StartWithNoopStoreClosesCleanly(t *testing.T) { + r := New(nil, nil, syncAuth, 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.go b/internal/contextcache/redis/store.go similarity index 58% rename from internal/contextcache/redis.go rename to internal/contextcache/redis/store.go index b62052336e..6443b1138d 100644 --- a/internal/contextcache/redis.go +++ b/internal/contextcache/redis/store.go @@ -3,7 +3,7 @@ // The full text of the Apache license is available in the LICENSE file at // the root of the repo. -package contextcache +package redis import ( "context" @@ -12,70 +12,72 @@ import ( "strings" "time" - "github.com/redis/go-redis/v9" + goredis "github.com/redis/go-redis/v9" + + "github.com/envoyproxy/ai-gateway/internal/contextcache" ) -// redisStore is a Store backed by Redis, shared across gateway replicas. +// redisStore is a contextcache.Store backed by Redis, shared across gateway replicas. type redisStore struct { - client redis.UniversalClient + client goredis.UniversalClient } -// NewRedisStore returns a ReconcileStore backed by the Redis instance at url, which accepts +// 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 NewRedisStore(url string) (ReconcileStore, error) { +func NewStore(url string) (contextcache.Store, error) { opts, err := parseRedisURL(url) if err != nil { return nil, err } - return &redisStore{client: redis.NewClient(opts)}, nil + return &redisStore{client: goredis.NewClient(opts)}, nil } // parseRedisURL accepts both a scheme-qualified URL and a bare host:port. -func parseRedisURL(url string) (*redis.Options, error) { +func parseRedisURL(url string) (*goredis.Options, error) { if url == "" { return nil, errors.New("contextcache: redis url is empty") } if strings.Contains(url, "://") { - opts, err := redis.ParseURL(url) + opts, err := goredis.ParseURL(url) if err != nil { return nil, fmt.Errorf("contextcache: invalid redis url: %w", err) } return opts, nil } - return &redis.Options{Addr: url}, nil + return &goredis.Options{Addr: url}, nil } // Get implements Store. -func (s *redisStore) Get(ctx context.Context, key string) (Entry, bool, error) { +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, redis.Nil) { - return Entry{}, false, nil + if errors.Is(err, goredis.Nil) { + return contextcache.Entry{}, false, nil } - return Entry{}, false, fmt.Errorf("contextcache: redis get: %w", err) + 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 Entry{}, false, nil + return contextcache.Entry{}, false, nil } return e, true, nil } // Set implements Store. -func (s *redisStore) Set(ctx context.Context, key string, e Entry, ttl time.Duration) error { +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 ReconcileStore. -func (s *redisStore) SetNX(ctx context.Context, key string, e Entry, ttl time.Duration) (bool, error) { +// 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) @@ -83,7 +85,7 @@ func (s *redisStore) SetNX(ctx context.Context, key string, e Entry, ttl time.Du return wrote, nil } -// AcquireGate implements ReconcileStore. +// 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 { @@ -92,22 +94,22 @@ func (s *redisStore) AcquireGate(ctx context.Context, key string, ttl time.Durat return won, nil } -var _ ReconcileStore = (*redisStore)(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 Entry) string { +func encodeEntry(e contextcache.Entry) string { return e.ExpireTime.UTC().Format(time.RFC3339) + "|" + e.Name } -func decodeEntry(s string) (Entry, error) { +func decodeEntry(s string) (contextcache.Entry, error) { expiry, name, ok := strings.Cut(s, "|") if !ok || name == "" { - return Entry{}, fmt.Errorf("contextcache: malformed cache entry %q", s) + return contextcache.Entry{}, fmt.Errorf("contextcache: malformed cache entry %q", s) } t, err := time.Parse(time.RFC3339, expiry) if err != nil { - return Entry{}, fmt.Errorf("contextcache: malformed cache entry expiry %q: %w", expiry, err) + return contextcache.Entry{}, fmt.Errorf("contextcache: malformed cache entry expiry %q: %w", expiry, err) } - return Entry{Name: name, ExpireTime: t}, nil + return contextcache.Entry{Name: name, ExpireTime: t}, nil } diff --git a/internal/contextcache/redis_test.go b/internal/contextcache/redis/store_test.go similarity index 89% rename from internal/contextcache/redis_test.go rename to internal/contextcache/redis/store_test.go index 7030f58dbb..7b1569e596 100644 --- a/internal/contextcache/redis_test.go +++ b/internal/contextcache/redis/store_test.go @@ -3,7 +3,7 @@ // The full text of the Apache license is available in the LICENSE file at // the root of the repo. -package contextcache +package redis import ( "context" @@ -13,13 +13,15 @@ import ( "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 := NewRedisStore(mr.Addr()) + s, err := NewStore(mr.Addr()) require.NoError(t, err) return s.(*redisStore), mr } @@ -51,7 +53,7 @@ func TestRedisStore_SetGetRoundTrip(t *testing.T) { ctx := context.Background() expire := time.Now().Add(10 * time.Minute).UTC().Truncate(time.Second) - require.NoError(t, s.Set(ctx, "k", Entry{Name: "projects/p/locations/r/cachedContents/x", ExpireTime: expire}, time.Minute)) + 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) @@ -81,7 +83,7 @@ func TestRedisStore_MalformedValueIsAMiss(t *testing.T) { 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 := Entry{Name: "projects/p|weird/cachedContents/x", ExpireTime: expire} + e := contextcache.Entry{Name: "projects/p|weird/cachedContents/x", ExpireTime: expire} got, err := decodeEntry(encodeEntry(e)) require.NoError(t, err) 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 index dd9d8db874..6e4b088a6c 100644 --- a/internal/contextcache/store.go +++ b/internal/contextcache/store.go @@ -5,10 +5,15 @@ // 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 -// reconciler that repairs drift between that store and the provider. +// 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 reconciler. +// 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 ( @@ -28,7 +33,7 @@ type Entry struct { // 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. It is a cache, not a lock: replicas that miss the +// 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. // @@ -40,13 +45,6 @@ type Store interface { 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 -} - -// ReconcileStore is implemented by stores that can back a Reconciler. It is kept out of -// Store so that the no-op store and test stores need not implement it; a store that does -// not implement it simply has no reconciler. -type ReconcileStore interface { - Store // 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. @@ -59,6 +57,9 @@ type ReconcileStore interface { // // 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. @@ -67,4 +68,14 @@ func (NoopStore) Get(context.Context, string) (Entry, bool, error) { return Entr // 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/reconciler.go b/internal/contextcache/syncer.go similarity index 74% rename from internal/contextcache/reconciler.go rename to internal/contextcache/syncer.go index 8183e71081..25343d09a9 100644 --- a/internal/contextcache/reconciler.go +++ b/internal/contextcache/syncer.go @@ -8,17 +8,28 @@ 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 reconciler does not write entries this close to expiry. +// 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 reconciliation. +// 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. @@ -29,7 +40,7 @@ type Source interface { List(ctx context.Context, add func(key string, e Entry)) error } -// Stats summarizes one reconcile round. +// Stats summarizes one sync round. type Stats struct { // Gated is true when another replica held the gate and nothing was listed. Gated bool @@ -42,21 +53,18 @@ type Stats struct { Skipped int } -// Reconciler repairs drift between a Store and a provider: caches that exist at the +// 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. // -// It never serves requests. If it is late, stuck, or stopped, requests still succeed; -// the only effect is that more of them create a cache instead of reusing one. -// // 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 Reconciler struct { - Store ReconcileStore +type Syncer struct { + Store Store Source Source Interval time.Duration Logger *slog.Logger @@ -64,7 +72,7 @@ type Reconciler struct { // Run runs rounds until ctx is done. It never returns an error: a failed round is logged // and the next one retries. -func (rc *Reconciler) Run(ctx context.Context) { +func (rc *Syncer) Run(ctx context.Context) { ticker := time.NewTicker(rc.Interval) defer ticker.Stop() for { @@ -77,17 +85,17 @@ func (rc *Reconciler) Run(ctx context.Context) { } } -func (rc *Reconciler) round(ctx context.Context) { +func (rc *Syncer) round(ctx context.Context) { stats, err := rc.RunOnce(ctx) if err != nil { if ctx.Err() == nil { - rc.logger().Warn("contextcache: reconcile round failed", + 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: reconcile round complete", + 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)) } @@ -95,11 +103,11 @@ func (rc *Reconciler) round(ctx context.Context) { // RunOnce performs one round. An error ends the round early; entries written before it // stay written. -func (rc *Reconciler) RunOnce(ctx context.Context) (Stats, error) { +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 reconcile gate: %w", err) + return stats, fmt.Errorf("acquire sync gate: %w", err) } if !won { stats.Gated = true @@ -116,7 +124,7 @@ func (rc *Reconciler) RunOnce(ctx context.Context) (Stats, error) { return stats, err } -func (rc *Reconciler) write(ctx context.Context, key string, e Entry) bool { +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 { @@ -124,14 +132,14 @@ func (rc *Reconciler) write(ctx context.Context, key string, e Entry) bool { } wrote, err := rc.Store.SetNX(ctx, key, e, ttl) if err != nil { - rc.logger().Warn("contextcache: reconcile store write failed", + rc.logger().Warn("contextcache: sync store write failed", slog.String("name", e.Name), slog.String("error", err.Error())) return false } return wrote } -func (rc *Reconciler) logger() *slog.Logger { +func (rc *Syncer) logger() *slog.Logger { if rc.Logger == nil { return slog.New(slog.DiscardHandler) } diff --git a/internal/contextcache/reconciler_test.go b/internal/contextcache/syncer_test.go similarity index 51% rename from internal/contextcache/reconciler_test.go rename to internal/contextcache/syncer_test.go index afa6ae69b9..ff7d34c667 100644 --- a/internal/contextcache/reconciler_test.go +++ b/internal/contextcache/syncer_test.go @@ -8,6 +8,7 @@ package contextcache import ( "context" "errors" + "sync" "testing" "time" @@ -33,15 +34,60 @@ func (s *fakeSource) List(_ context.Context, add func(string, Entry)) error { return s.err } -func newTestReconciler(t *testing.T, src *fakeSource) (*Reconciler, ReconcileStore) { - t.Helper() - store, _ := newTestRedisStore(t) - return &Reconciler{Store: store, Source: src, Interval: time.Minute}, store +// 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 TestReconciler_WritesEntries(t *testing.T) { +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 := newTestReconciler(t, &fakeSource{gate: "g", entries: map[string]Entry{ + rc, store := newTestSyncer(t, &fakeSource{gate: "g", entries: map[string]Entry{ "a": {Name: "c/a", ExpireTime: exp}, "b": {Name: "c/b", ExpireTime: exp}, }}) @@ -55,9 +101,9 @@ func TestReconciler_WritesEntries(t *testing.T) { assert.Equal(t, "c/a", e.Name) } -func TestReconciler_DoesNotOverwrite(t *testing.T) { +func TestSyncer_DoesNotOverwrite(t *testing.T) { exp := time.Now().Add(5 * time.Minute) - rc, store := newTestReconciler(t, &fakeSource{gate: "g", entries: map[string]Entry{ + 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)) @@ -69,8 +115,8 @@ func TestReconciler_DoesNotOverwrite(t *testing.T) { assert.Equal(t, "c/published", e.Name) } -func TestReconciler_SkipsNearExpiry(t *testing.T) { - rc, store := newTestReconciler(t, &fakeSource{gate: "g", entries: map[string]Entry{ +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)}, }}) @@ -81,9 +127,9 @@ func TestReconciler_SkipsNearExpiry(t *testing.T) { assert.False(t, ok) } -func TestReconciler_GateSkipsSecondRound(t *testing.T) { +func TestSyncer_GateSkipsSecondRound(t *testing.T) { src := &fakeSource{gate: "g"} - rc, _ := newTestReconciler(t, src) + rc, _ := newTestSyncer(t, src) _, err := rc.RunOnce(context.Background()) require.NoError(t, err) @@ -93,9 +139,9 @@ func TestReconciler_GateSkipsSecondRound(t *testing.T) { assert.Equal(t, 1, src.calls, "a gated round must not list") } -func TestReconciler_ListErrorKeepsEarlierWrites(t *testing.T) { +func TestSyncer_ListErrorKeepsEarlierWrites(t *testing.T) { exp := time.Now().Add(5 * time.Minute) - rc, store := newTestReconciler(t, &fakeSource{ + rc, store := newTestSyncer(t, &fakeSource{ gate: "g", err: errors.New("boom"), entries: map[string]Entry{"a": {Name: "c/a", ExpireTime: exp}}, }) @@ -106,8 +152,8 @@ func TestReconciler_ListErrorKeepsEarlierWrites(t *testing.T) { assert.True(t, ok) } -func TestReconciler_RunStopsOnCancel(t *testing.T) { - rc, _ := newTestReconciler(t, &fakeSource{gate: "g", err: errors.New("boom")}) +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{}) @@ -122,3 +168,19 @@ func TestReconciler_RunStopsOnCancel(t *testing.T) { 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/extproc/processor_impl.go b/internal/extproc/processor_impl.go index afbb35557f..5ce8341d7f 100644 --- a/internal/extproc/processor_impl.go +++ b/internal/extproc/processor_impl.go @@ -27,9 +27,9 @@ import ( "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/gcpcache" "github.com/envoyproxy/ai-gateway/internal/headermutator" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/json" @@ -130,9 +130,9 @@ type ( backendName string routeName string handler filterapi.BackendAuthHandler - // cacheResolver is the GCP context-cache resolver for this backend. + // cacheResolver is context-cache resolver for this backend. // Non-nil only for GCP Vertex AI backends; set in SetBackend. - cacheResolver gcpcache.CacheResolver + 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. @@ -373,30 +373,34 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque // We force the body mutation in the following cases: // * The request is a retry request because the body mutation might have happened the previous iteration. // * 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 pypassed. + // 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. if u.cacheResolver != nil { - if cacheSetter, ok := u.translator.(translator.GCPCacheSetter); ok { - if gcpAuth, ok := u.handler.(filterapi.GCPAuthHandler); ok { - if req, ok := any(u.parent.originalRequestBody).(*openai.ChatCompletionRequest); ok { - result, resolveErr := u.cacheResolver.Resolve(ctx, req, gcpAuth) - if resolveErr != nil { - u.logger.Error("GCP 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 cacheSetter, ok := u.translator.(translator.ContextCacheSetter); ok { + // Skip requests that name their own cache. + if req, ok := any(u.parent.originalRequestBody).(*openai.ChatCompletionRequest); ok && + (req.GCPVertexAIVendorFields == nil || req.CachedContent == "") { + 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, } - if result != nil { - cacheSetter.SetGCPCacheResult(&translator.GCPCacheResult{ - CacheName: result.CacheName, - FilteredMessages: result.Messages, - Created: result.Created, - WriteTokenCount: uint32(result.TokenCount), //nolint:gosec - }) + // 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) } } } @@ -746,9 +750,7 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) SetBackend(c u.backendName = backend.Backend.Name u.routeName = routeName u.handler = backend.Handler - if cr, ok := backend.CacheResolver.(gcpcache.CacheResolver); ok { - u.cacheResolver = cr - } + 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 7a45b689c6..5729af26ed 100644 --- a/internal/extproc/processor_impl_test.go +++ b/internal/extproc/processor_impl_test.go @@ -31,9 +31,9 @@ import ( "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/gcpcache" "github.com/envoyproxy/ai-gateway/internal/headermutator" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/json" @@ -2525,25 +2525,25 @@ func mustCompileCEL(t *testing.T, expr string) cel.Program { // GCP context-cache resolver wiring tests // --------------------------------------------------------------------------- -// mockGCPCacheResolver is a controllable gcpcache.CacheResolver for testing. -type mockGCPCacheResolver struct { - retResult *gcpcache.ResolveResult +// mockCacheResolver is a controllable contextcache.CacheResolver for testing. +type mockCacheResolver struct { + retResult *contextcache.ResolveResult retErr error called bool } -func (m *mockGCPCacheResolver) Resolve(_ context.Context, _ *openai.ChatCompletionRequest, _ filterapi.GCPAuthHandler) (*gcpcache.ResolveResult, error) { +func (m *mockCacheResolver) Resolve(_ context.Context, _ *openai.ChatCompletionRequest) (*contextcache.ResolveResult, error) { m.called = true return m.retResult, m.retErr } -// mockGCPCacheTranslator extends mockTranslator with GCPCacheSetter support. -type mockGCPCacheTranslator struct { +// mockCacheTranslator extends mockTranslator with ContextCacheSetter support. +type mockCacheTranslator struct { mockTranslator - setCacheResult *translator.GCPCacheResult + setCacheResult *translator.ContextCacheResult } -func (m *mockGCPCacheTranslator) SetGCPCacheResult(result *translator.GCPCacheResult) { +func (m *mockCacheTranslator) SetContextCacheResult(result *translator.ContextCacheResult) { m.setCacheResult = result } @@ -2561,15 +2561,15 @@ func (m *mockGCPAuthHandler) GCPTokenSource() oauth2.TokenSource { func (m *mockGCPAuthHandler) GCPRegion() string { return "us-central1" } func (m *mockGCPAuthHandler) GCPProject() string { return "test-project" } -// Test_upstreamProcessor_GCPCacheResolver_notCalled verifies that the resolver is not invoked -// when the translator does not implement GCPCacheSetter. -func Test_upstreamProcessor_GCPCacheResolver_notCalled(t *testing.T) { +// 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 := &mockGCPCacheResolver{} - // mockTranslator does not implement GCPCacheSetter. + cr := &mockCacheResolver{} + // mockTranslator does not implement ContextCacheSetter. tr := &mockTranslator{t: t, expRequestBody: &body} mm := &mockMetrics{} @@ -2589,18 +2589,18 @@ func Test_upstreamProcessor_GCPCacheResolver_notCalled(t *testing.T) { } _, err := p.ProcessRequestHeaders(t.Context(), nil) require.NoError(t, err) - require.False(t, cr.called, "resolver must not be called when translator lacks GCPCacheSetter") + require.False(t, cr.called, "resolver must not be called when translator lacks ContextCacheSetter") } -// Test_upstreamProcessor_GCPCacheResolver_noMarkers verifies that when the resolver returns nil +// 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_GCPCacheResolver_noMarkers(t *testing.T) { +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 := &mockGCPCacheResolver{retResult: nil} // no markers - tr := &mockGCPCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + cr := &mockCacheResolver{retResult: nil} // no markers + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} mm := &mockMetrics{} p := &chatCompletionProcessorUpstreamFilter{ @@ -2620,22 +2620,22 @@ func Test_upstreamProcessor_GCPCacheResolver_noMarkers(t *testing.T) { _, err := p.ProcessRequestHeaders(t.Context(), nil) require.NoError(t, err) require.True(t, cr.called) - require.Nil(t, tr.setCacheResult, "SetGCPCacheResult must not be called when resolver returns nil") + require.Nil(t, tr.setCacheResult, "SetContextCacheResult must not be called when resolver returns nil") } -// Test_upstreamProcessor_GCPCacheResolver_hit verifies that a cache hit is injected into the translator. -func Test_upstreamProcessor_GCPCacheResolver_hit(t *testing.T) { +// 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 := &mockGCPCacheResolver{retResult: &gcpcache.ResolveResult{ + cr := &mockCacheResolver{retResult: &contextcache.ResolveResult{ CacheName: "projects/p/locations/us-central1/cachedContents/abc123", Messages: filteredMessages, Created: false, }} - tr := &mockGCPCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} mm := &mockMetrics{} p := &chatCompletionProcessorUpstreamFilter{ @@ -2662,19 +2662,19 @@ func Test_upstreamProcessor_GCPCacheResolver_hit(t *testing.T) { require.Zero(t, tr.setCacheResult.WriteTokenCount) } -// Test_upstreamProcessor_GCPCacheResolver_created verifies that cache-write tokens are propagated. -func Test_upstreamProcessor_GCPCacheResolver_created(t *testing.T) { +// 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 := &mockGCPCacheResolver{retResult: &gcpcache.ResolveResult{ + cr := &mockCacheResolver{retResult: &contextcache.ResolveResult{ CacheName: "projects/p/locations/us-central1/cachedContents/new456", Messages: []openai.ChatCompletionMessageParamUnion{}, Created: true, TokenCount: 512, }} - tr := &mockGCPCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + tr := &mockCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} mm := &mockMetrics{} p := &chatCompletionProcessorUpstreamFilter{ @@ -2699,50 +2699,60 @@ func Test_upstreamProcessor_GCPCacheResolver_created(t *testing.T) { require.Equal(t, uint32(512), tr.setCacheResult.WriteTokenCount) } -// Test_upstreamProcessor_GCPCacheResolver_error verifies that a resolver error produces a 502. -func Test_upstreamProcessor_GCPCacheResolver_error(t *testing.T) { +// 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 := &mockGCPCacheResolver{retErr: errors.New("GCP API quota exceeded")} - tr := &mockGCPCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} - mm := &mockMetrics{} - + 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", + 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(), + metrics: &mockMetrics{}, translator: tr, handler: &mockGCPAuthHandler{}, cacheResolver: cr, } - resp, err := p.ProcessRequestHeaders(t.Context(), nil) + _, 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) + require.NotNil(t, tr.setCacheResult) + require.Zero(t, tr.setCacheResult.WriteTokenCount) } -// Test_upstreamProcessor_GCPCacheResolver_noHandler verifies the resolver is skipped -// when the handler is not a GCPAuthHandler. -func Test_upstreamProcessor_GCPCacheResolver_noHandler(t *testing.T) { +// A request that names its own cache must not be resolved, so no cache is created for it. +func Test_upstreamProcessor_CacheResolver_explicitCachedContentSkips(t *testing.T) { someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) var body openai.ChatCompletionRequest require.NoError(t, json.Unmarshal(someBody, &body)) + body.GCPVertexAIVendorFields = &openai.GCPVertexAIVendorFields{CachedContent: "c/explicit"} - cr := &mockGCPCacheResolver{} - tr := &mockGCPCacheTranslator{mockTranslator: mockTranslator{t: t, expRequestBody: &body}} + cr := &mockCacheResolver{} + 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, + } + _, _ = p.ProcessRequestHeaders(t.Context(), nil) + require.False(t, cr.called) + require.Nil(t, tr.setCacheResult) +} + +// 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{ @@ -2756,11 +2766,16 @@ func Test_upstreamProcessor_GCPCacheResolver_noHandler(t *testing.T) { requestHeaders: map[string]string{internalapi.ModelNameHeaderKeyDefault: "gemini-1.5-pro"}, metrics: mm, translator: tr, - handler: &mockBackendAuthHandler{}, // not a GCPAuthHandler + handler: &mockGCPAuthHandler{}, cacheResolver: cr, + logger: slog.Default(), } - _, err := p.ProcessRequestHeaders(t.Context(), nil) + resp, err := p.ProcessRequestHeaders(t.Context(), nil) require.NoError(t, err) - require.False(t, cr.called, "resolver must not be called when handler is not GCPAuthHandler") - require.Nil(t, tr.setCacheResult) + 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 bb2c0e8401..1f1d2fc441 100644 --- a/internal/extproc/server.go +++ b/internal/extproc/server.go @@ -26,9 +26,9 @@ 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/contextcache/redis" "github.com/envoyproxy/ai-gateway/internal/filterapi" - "github.com/envoyproxy/ai-gateway/internal/gcpcache" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/redaction" ) @@ -68,7 +68,9 @@ type Server struct { // filterapi update does not tear down and rebuild the Redis connection pool. The // Redis URL is part of the key so that repointing a backend at a different Redis // builds a new resolver rather than reusing a pool aimed at the old one. - cacheResolvers map[string]gcpcache.CacheResolver + cacheResolvers map[string]*gcpcache.Resolver + // cacheResolversMu guards cacheResolvers against LoadConfig racing Close. + cacheResolversMu sync.Mutex } // NewServer creates a new external processor server. @@ -80,7 +82,7 @@ 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.CacheResolver), + cacheResolvers: make(map[string]*gcpcache.Resolver), uuidFn: uuid.NewString, } return srv, nil @@ -100,11 +102,15 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error logger = slog.New(slog.DiscardHandler) } - // Attach a context-cache resolver to each GCP Vertex AI backend that configures - // context caching. Resolvers are reused across reloads so that a config update does - // not tear down and rebuild each backend's Redis connection pool; resolvers for - // backends that disappeared or dropped their cache config are discarded. - live := make(map[string]gcpcache.CacheResolver, len(s.cacheResolvers)) + // Attach a context-cache resolver to each backend with GCP credentials that + // configures context caching. Resolvers are keyed by backend name and Redis URL and + // reused across reloads, so a config update does not rebuild their Redis connection + // pools; a reused resolver is given the reload's auth handler, since the controller + // rotates the token. Resolvers whose backend disappeared, dropped its cache config, or + // moved to another Redis 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 _, ok := rb.Handler.(filterapi.GCPAuthHandler); !ok { continue @@ -116,7 +122,7 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error key := rb.Backend.Name + "|" + cc.URL resolver, ok := s.cacheResolvers[key] if !ok { - store, err := contextcache.NewRedisStore(cc.URL) + store, err := redis.NewStore(cc.URL) if err != nil { // A bad URL disables caching for this backend rather than failing the // reload: the store is an optimization, and rejecting the whole config @@ -125,14 +131,16 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error slog.String("backend", rb.Backend.Name), slog.String("error", err.Error())) continue } - resolver = gcpcache.New(nil, store, logger) - } - if bg, ok := resolver.(gcpcache.Background); ok { - // Start is a no-op for a resolver already running, so reused resolvers are - // unaffected. It runs under the resolver's own context, not ctx, which is - // cancelled when this reload finishes. - bg.Start(rb.Handler.(filterapi.GCPAuthHandler)) - } + resolver = gcpcache.New(nil, store, rb.Handler.(filterapi.GCPAuthHandler), logger) + } else { + // The handler is rebuilt on every reload (the controller rotates the access + // token), so hand the reused resolver the current one. + resolver.SetAuth(rb.Handler.(filterapi.GCPAuthHandler)) + } + // Start is a no-op for a resolver already running, so reused resolvers are + // unaffected. It runs under the resolver's own context, not ctx, which is + // cancelled when this reload finishes. + resolver.Start() live[key] = resolver rb.CacheResolver = resolver } @@ -142,9 +150,7 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error if _, ok := live[key]; ok { continue } - if c, ok := resolver.(io.Closer); ok { - _ = c.Close() - } + _ = resolver.Close() } s.cacheResolvers = live @@ -152,13 +158,13 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error return nil } -// Close stops the background work of every context-cache resolver. It is safe to call +// 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 { - if c, ok := resolver.(io.Closer); ok { - _ = c.Close() - } + _ = resolver.Close() } } diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index b144768b81..bd97a37a69 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -137,6 +137,23 @@ func TestServer_LoadConfig_GCPCacheResolver_ReusedAcrossReloads(t *testing.T) { "an unchanged backend must keep its resolver, and with it its connection pool") } +// A reused resolver must pick up the handler rebuilt on reload, since the controller +// rotates the access token without changing the Redis URL. +func TestServer_LoadConfig_GCPCacheResolver_ReusedResolverGetsNewHandler(t *testing.T) { + s := &Server{} + t.Cleanup(s.Close) + config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}} + + require.NoError(t, s.LoadConfig(t.Context(), config)) + first := s.config.Backends["gcp"].Handler + + require.NoError(t, s.LoadConfig(t.Context(), config)) + second := s.config.Backends["gcp"].Handler + require.NotSame(t, first, second, "the handler is rebuilt on every reload") + require.Equal(t, second, s.cacheResolvers["gcp|localhost:6379"].GetAuth(), + "the reused resolver must hold the current handler") +} + // Repointing a backend at a different Redis must not keep reusing a pool dialing the old one. func TestServer_LoadConfig_GCPCacheResolver_RebuiltWhenRedisURLChanges(t *testing.T) { s := &Server{} @@ -187,7 +204,7 @@ func TestServer_LoadConfig_GCPCacheResolver_BadRedisURLDoesNotFailReload(t *test "an unrelated backend must keep its resolver") } -// A resolver evicted by a reload must have its background reconciler stopped, or every +// A resolver evicted by a reload must have its background syncer stopped, or every // reload that drops or repoints a backend would leak a goroutine. func TestServer_LoadConfig_GCPCacheResolver_EvictedResolverIsClosed(t *testing.T) { before := goleak.IgnoreCurrent() @@ -197,20 +214,20 @@ func TestServer_LoadConfig_GCPCacheResolver_EvictedResolverIsClosed(t *testing.T gcpBackendWithCaching("gcp-a", "localhost:6379"), gcpBackendWithCaching("gcp-b", "localhost:6379"), }})) - kept := s.cacheResolvers["gcp-a|localhost:6379"].(io.Closer) + kept := s.cacheResolvers["gcp-a|localhost:6379"] require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ gcpBackendWithCaching("gcp-a", "localhost:6379"), }})) // Closing the one resolver still in the config must leave nothing running. If the - // evicted resolver had not been closed by the reload, its reconciler would remain. + // evicted resolver had not been closed by the reload, its syncer would remain. require.NoError(t, kept.Close()) goleak.VerifyNone(t, before) } -// The reload context is cancelled when the reload finishes. Reconcilers must not be tied +// The reload context is cancelled when the reload finishes. Syncers must not be tied // to it, or they would stop seconds after every config load. -func TestServer_LoadConfig_ReconcilerOutlivesReloadContext(t *testing.T) { +func TestServer_LoadConfig_SyncerOutlivesReloadContext(t *testing.T) { before := goleak.IgnoreCurrent() s := &Server{} @@ -220,7 +237,7 @@ func TestServer_LoadConfig_ReconcilerOutlivesReloadContext(t *testing.T) { }})) cancel() time.Sleep(50 * time.Millisecond) - require.Error(t, goleak.Find(before), "the reconciler must still be running after the reload context ends") + require.Error(t, goleak.Find(before), "the syncer must still be running after the reload context ends") s.Close() goleak.VerifyNone(t, before) diff --git a/internal/filterapi/filterconfig.go b/internal/filterapi/filterconfig.go index 1aec42eda9..f82d151458 100644 --- a/internal/filterapi/filterconfig.go +++ b/internal/filterapi/filterconfig.go @@ -208,7 +208,7 @@ type Backend struct { // the request to the backend. Optional. HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` // ContextCache configures context caching for this backend. Optional. - // When non-nil, the extproc will resolve or create a Google cachedContents entry + // When non-nil, the extproc will resolve or create a GCP cachedContents entry // for requests that carry cache_control markers. ContextCache *ContextCache `json:"contextCache,omitempty"` } @@ -219,12 +219,12 @@ type Backend struct { // cache entry is created when absent. type ContextCache struct { // DefaultTTL is the default time-to-live for newly created cache entries expressed - // as a GCP duration string (e.g. "600s"). When empty, the resolver default (300s) is used. + // as a duration string (e.g. "600s"). When empty, the resolver default (300s) is used. // +optional DefaultTTL string `json:"defaultTTL,omitempty"` // URL points at the shared Redis store that records which cachedContents entry a - // given prefix resolved to. It is what stops two gateway replicas from - // independently creating a cache for the same prefix. + // given prefix resolved to. Once one replica has created a cache, the others reuse + // it. Replicas that miss the same prefix at the same moment may each create one. // // It is either a bare "host:port" or a full "redis://" URL. URL string `json:"url"` diff --git a/internal/filterapi/runtime.go b/internal/filterapi/runtime.go index caa061e00b..695a8ee54d 100644 --- a/internal/filterapi/runtime.go +++ b/internal/filterapi/runtime.go @@ -12,6 +12,7 @@ import ( "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" ) @@ -68,11 +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. - // It is populated by the extproc layer for GCP Vertex AI backends (where Handler - // implements GCPAuthHandler). Held as any to avoid an import cycle between - // filterapi and gcpcache; callers type-assert to gcpcache.CacheResolver. - CacheResolver any + // 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/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index 37f0f10c52..2be22614c6 100644 --- a/internal/translator/openai_gcpvertexai.go +++ b/internal/translator/openai_gcpvertexai.go @@ -92,12 +92,12 @@ type openAIToGCPVertexAITranslatorV1ChatCompletion struct { debugLogEnabled bool enableRedaction bool logger *slog.Logger - // pendingCacheResult is set by SetGCPCacheResult before RequestBody is called. + // 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. - pendingCacheResult *GCPCacheResult + cacheResult *ContextCacheResult // cacheWriteTokens is the number of tokens stored in a newly created cache entry. - // Set during RequestBody (from pendingCacheResult) and consumed in ResponseBody to + // Set during RequestBody (from cacheResult) and consumed in ResponseBody to // record cache-write cost via tokenUsage.SetCacheCreationInputTokens. cacheWriteTokens uint32 } @@ -135,11 +135,10 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) RequestBody(_ []byte, op return nil, nil, err } - // Capture cache-write token count for cost attribution in ResponseBody, then clear - // the pending result so retries start fresh (the resolver will re-check the memo). - if o.pendingCacheResult != nil { - o.cacheWriteTokens = o.pendingCacheResult.WriteTokenCount - o.pendingCacheResult = nil + // 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) @@ -547,11 +546,11 @@ 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 is pending (set by SetGCPCacheResult), use the filtered (non-cached) + // 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.pendingCacheResult != nil && o.pendingCacheResult.FilteredMessages != nil { - messages = o.pendingCacheResult.FilteredMessages + if o.cacheResult != nil && o.cacheResult.FilteredMessages != nil { + messages = o.cacheResult.FilteredMessages } // Convert OpenAI messages to Gemini Contents and SystemInstruction. @@ -596,11 +595,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 SetGCPCacheResult). + // Inject resolved cache name from the context-cache resolver (set via SetContextCacheResult). // This happens after vendor-field processing so an explicit vendor cachedContent still wins // (the validation in RequestBody already rejected the mixed case). - if o.pendingCacheResult != nil && gcr.CachedContent == "" { - gcr.CachedContent = o.pendingCacheResult.CacheName + if o.cacheResult != nil && gcr.CachedContent == "" { + 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 @@ -675,12 +679,12 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) applyVendorSpecificField } } -// SetGCPCacheResult implements [GCPCacheSetter]. The upstream processor calls this +// 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) SetGCPCacheResult(result *GCPCacheResult) { - o.pendingCacheResult = result +func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) SetContextCacheResult(result *ContextCacheResult) { + o.cacheResult = result } func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) geminiResponseToOpenAIMessage(gcr *genai.GenerateContentResponse, responseModel string) (*openai.ChatCompletionResponse, error) { diff --git a/internal/translator/openai_gcpvertexai_test.go b/internal/translator/openai_gcpvertexai_test.go index bb7b235308..60896bbbf4 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -2922,15 +2922,15 @@ func TestGCPVertexAIRedactBody(t *testing.T) { }) } -// TestOpenAIToGCPVertexAITranslator_GCPCacheSetter verifies that SetGCPCacheResult +// 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_GCPCacheSetter(t *testing.T) { +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.SetGCPCacheResult(&GCPCacheResult{ + tr.SetContextCacheResult(&ContextCacheResult{ CacheName: "projects/p/locations/us-central1/cachedContents/abc", FilteredMessages: []openai.ChatCompletionMessageParamUnion{}, Created: true, @@ -2953,16 +2953,47 @@ func TestOpenAIToGCPVertexAITranslator_GCPCacheSetter(t *testing.T) { // The Gemini request should carry the cachedContent field. require.Contains(t, string(bodyBytes), `"cachedContent":"projects/p/locations/us-central1/cachedContents/abc"`) - // pendingCacheResult should be cleared after RequestBody. - require.Nil(t, tr.pendingCacheResult) // 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.SetGCPCacheResult(&GCPCacheResult{ + tr.SetContextCacheResult(&ContextCacheResult{ CacheName: "projects/p/locations/us-central1/cachedContents/existing", Created: false, WriteTokenCount: 0, @@ -2977,6 +3008,14 @@ func TestOpenAIToGCPVertexAITranslator_GCPCacheSetter(t *testing.T) { 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 diff --git a/internal/translator/translator.go b/internal/translator/translator.go index 4fcdc60571..e3c36007c1 100644 --- a/internal/translator/translator.go +++ b/internal/translator/translator.go @@ -117,26 +117,27 @@ type ResponseRedactor interface { RedactBody(resp *openai.ChatCompletionResponse) *openai.ChatCompletionResponse } -// GCPCacheSetter is an optional interface implemented by the GCP Vertex AI translator. -// The upstream processor calls SetGCPCacheResult after the context-cache resolver runs so -// that the translator can (a) inject the cache resource name into the Gemini request and -// (b) record cache-write token counts for cost attribution in the response phase. -type GCPCacheSetter interface { - // SetGCPCacheResult forwards the resolved (or newly created) cache entry to the translator. +// 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. - SetGCPCacheResult(result *GCPCacheResult) + SetContextCacheResult(result *ContextCacheResult) } -// GCPCacheResult carries the information returned by the cache resolver that the translator needs. -type GCPCacheResult struct { - // CacheName is the full Google resource name of the resolved or created cache entry. +// 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 Gemini request. + // 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 Google's create response). + // WriteTokenCount is the number of tokens stored in the cache, from the provider's create response. // Only meaningful when Created is true. WriteTokenCount uint32 } diff --git a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml index ef7342939b..296007ed91 100644 --- a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml +++ b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml @@ -242,8 +242,7 @@ spec: url: description: |- URL locates the shared Redis store that records which cachedContents entry a given - request prefix resolved to. Sharing that state across gateway replicas is what - prevents two replicas from independently creating a cache for the same prefix. + request prefix resolved to. It is given either as a bare "host:port" (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. @@ -690,8 +689,7 @@ spec: url: description: |- URL locates the shared Redis store that records which cachedContents entry a given - request prefix resolved to. Sharing that state across gateway replicas is what - prevents two replicas from independently creating a cache for the same prefix. + request prefix resolved to. It is given either as a bare "host:port" (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. diff --git a/site/docs/api/api.mdx b/site/docs/api/api.mdx index d83040871f..1139e574f9 100644 --- a/site/docs/api/api.mdx +++ b/site/docs/api/api.mdx @@ -1413,7 +1413,7 @@ Leaving it unset disables caching entirely, and markers are ignored. name="url" type="string" required="true" - description="URL locates the shared Redis store that records which cachedContents entry a given
request prefix resolved to. Sharing that state across gateway replicas is what
prevents two replicas from independently creating a cache for the same prefix.
It is given either as a bare `host:port`
(e.g. `redis.default.svc.cluster.local:6379`) or as a full `redis://` URL.
If Redis is unreachable at request time, caching is skipped and the request is
served uncached rather than failed." + description="URL locates the shared Redis store that records which cachedContents entry a given
request prefix resolved to.
It is given either as a bare `host:port`
(e.g. `redis.default.svc.cluster.local:6379`) or as a full `redis://` URL.
If Redis is unreachable at request time, caching is skipped and the request is
served uncached rather than failed." /> @@ -4072,7 +4072,7 @@ Leaving it unset disables caching entirely, and markers are ignored. name="url" type="string" required="true" - description="URL locates the shared Redis store that records which cachedContents entry a given
request prefix resolved to. Sharing that state across gateway replicas is what
prevents two replicas from independently creating a cache for the same prefix.
It is given either as a bare `host:port`
(e.g. `redis.default.svc.cluster.local:6379`) or as a full `redis://` URL.
If Redis is unreachable at request time, caching is skipped and the request is
served uncached rather than failed." + description="URL locates the shared Redis store that records which cachedContents entry a given
request prefix resolved to.
It is given either as a bare `host:port`
(e.g. `redis.default.svc.cluster.local:6379`) or as a full `redis://` URL.
If Redis is unreachable at request time, caching is skipped and the request is
served uncached rather than failed." /> diff --git a/tests/crdcel/main_test.go b/tests/crdcel/main_test.go index 3466456ef5..902875cbda 100644 --- a/tests/crdcel/main_test.go +++ b/tests/crdcel/main_test.go @@ -95,10 +95,16 @@ func TestAIServiceBackends(t *testing.T) { }{ {name: "basic.yaml"}, {name: "anthropic-schema.yaml"}, - {name: "gcp-context-caching-enabled.yaml"}, + {name: "context-cache-enabled.yaml"}, { - name: "gcp-context-caching-invalid-ttl.yaml", - expErr: "spec.gcpContextCaching.defaultTTL", + name: "context-cache-invalid-ttl.yaml", + expErr: "spec.contextCache.defaultTTL", + }, + { + // Url should not be empty, so a missing + // url reaches the API server as "" and fails the MinLength rule. + name: "context-cache-missing-url.yaml", + expErr: "spec.contextCache.url in body should be at least 1 chars long", }, {name: "basic-eg-backend-aws.yaml"}, {name: "aws-openai-schema.yaml"}, diff --git a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml similarity index 87% rename from tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml rename to tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml index 745500215e..db43d02160 100644 --- a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-enabled.yaml +++ b/tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml @@ -16,6 +16,6 @@ spec: kind: Backend group: gateway.envoyproxy.io port: 443 - gcpContextCaching: - enabled: true + contextCache: + url: "redis.default.svc.cluster.local:6379" defaultTTL: "600s" diff --git a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml similarity index 88% rename from tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml rename to tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml index dad194c47c..4744640280 100644 --- a/tests/crdcel/testdata/aiservicebackends/gcp-context-caching-invalid-ttl.yaml +++ b/tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml @@ -16,6 +16,6 @@ spec: kind: Backend group: gateway.envoyproxy.io port: 443 - gcpContextCaching: - enabled: true + contextCache: + url: "redis.default.svc.cluster.local:6379" defaultTTL: "not-a-duration" diff --git a/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml new file mode 100644 index 0000000000..26edc20522 --- /dev/null +++ b/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml @@ -0,0 +1,20 @@ +# 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. + +apiVersion: aigateway.envoyproxy.io/v1beta1 +kind: AIServiceBackend +metadata: + name: gcp-cache-backend-no-url + namespace: default +spec: + schema: + name: GCPVertexAI + backendRef: + name: gcp-service + kind: Backend + group: gateway.envoyproxy.io + port: 443 + contextCache: + defaultTTL: "600s" From 9f9f525b04f17ebf72ad432e0491878960b8d7c3 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Tue, 6 Oct 2026 21:40:25 +0530 Subject: [PATCH 13/17] contextcache: configure the store and default TTL on extproc, not per backend Context caching was configured per backend through a contextCache field on AIServiceBackend (url, defaultTTL). Configure it once per gateway instead. This removes contextCache from AIServiceBackend in both v1alpha1 and v1beta1, and from filterapi.Backend, along with the controller code that copied it. Manifests that set it will have the field dropped. New configuration: - extproc flags -contextCacheStoreURL and -contextCacheDefaultTTL. - controller flags --extProcContextCacheStoreURL and --extProcContextCacheDefaultTTL, passed to extproc only when the store URL is set and added to the config hash with omitempty, so deployments without caching keep the same args and are not restarted on upgrade. - Helm values extProc.contextCache.storeURL and defaultTTL. The store URL is host:port or a redis:// URL of a Redis-compatible server. Setting it enables caching for every GCP Vertex AI backend with GCP credentials; there is no per-backend opt-out. The default TTL must be whole seconds and at least 60s, and is only validated when caching is on. extproc builds one store at startup and shares it across resolvers. A store URL that does not parse now fails startup instead of disabling caching for one backend. Resolvers are keyed by backend name and are only attached to GCPVertexAI backends, so GCPAnthropic backends no longer get a syncer that lists caches they cannot use. The default TTL is now passed to gcpcache.New. The old CRD defaultTTL was never read; the resolver always used a hardcoded 300s. Signed-off-by: Sivanantham Chinnaiyan --- api/v1alpha1/ai_service_backend.go | 36 ---- api/v1alpha1/zz_generated.deepcopy.go | 20 -- api/v1beta1/ai_service_backend.go | 36 ---- api/v1beta1/zz_generated.deepcopy.go | 20 -- cmd/controller/main.go | 16 ++ cmd/extproc/mainlib/main.go | 22 ++ cmd/extproc/mainlib/main_test.go | 16 ++ internal/contextcache/gcpcache/gcpcache.go | 33 ++- .../contextcache/gcpcache/gcpcache_test.go | 27 ++- internal/contextcache/gcpcache/syncer_test.go | 2 +- internal/controller/controller.go | 6 + internal/controller/extproc_builder.go | 31 ++- internal/controller/extproc_builder_test.go | 26 ++- internal/controller/gateway.go | 7 - internal/controller/gateway_test.go | 117 ---------- internal/extproc/server.go | 78 ++++--- internal/extproc/server_test.go | 200 +++++------------- internal/filterapi/filterconfig.go | 21 -- ...teway.envoyproxy.io_aiservicebackends.yaml | 60 ------ .../ai-gateway-helm/templates/deployment.yaml | 6 + manifests/charts/ai-gateway-helm/values.yaml | 9 + site/docs/api/api.mdx | 72 ------- tests/crdcel/main_test.go | 11 - .../context-cache-enabled.yaml | 21 -- .../context-cache-invalid-ttl.yaml | 21 -- .../context-cache-missing-url.yaml | 20 -- 26 files changed, 269 insertions(+), 665 deletions(-) delete mode 100644 tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml delete mode 100644 tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml delete mode 100644 tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml diff --git a/api/v1alpha1/ai_service_backend.go b/api/v1alpha1/ai_service_backend.go index c586e99a81..cddcc253f9 100644 --- a/api/v1alpha1/ai_service_backend.go +++ b/api/v1alpha1/ai_service_backend.go @@ -76,46 +76,10 @@ type AIServiceBackendSpec struct { // +optional BodyMutation *HTTPBodyMutation `json:"bodyMutation,omitempty"` - // ContextCache configures context caching for this backend. - // When set, the ai-gateway will automatically resolve or create a context cache - // entry for requests that carry Anthropic-style cache_control markers, and inject the - // resolved cache resource name into the API request. - // Currently only supported for GCP Vertex AI backends. - // +optional - ContextCache *ContextCacheSpec `json:"contextCache,omitempty"` - // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } -// ContextCacheSpec configures context caching for a backend. -// -// Setting this field enables context caching: cache_control markers in requests are -// resolved against the shared store at URL, and a cache entry is created when absent. -// Leaving it unset disables caching entirely, and markers are ignored. -type ContextCacheSpec struct { - // DefaultTTL is the default time-to-live for newly created cache entries, expressed - // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. - // Must be at least "60s" per GCP API requirements. - // - // +optional - // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` - DefaultTTL string `json:"defaultTTL,omitempty"` - - // URL locates the shared Redis store that records which cachedContents entry a given - // request prefix resolved to. - // - // It is given either as a bare "host:port" - // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. - // - // If Redis is unreachable at request time, caching is skipped and the request is - // served uncached rather than failed. - // - // +kubebuilder:validation:Required - // +kubebuilder:validation:MinLength=1 - URL string `json:"url"` -} - // HTTPHeaderMutation defines the mutation of HTTP headers that will be applied to the request type HTTPHeaderMutation struct { // Set overwrites/adds the request with the given header (name, value) diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 23d6ed6da9..f795d34163 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -338,11 +338,6 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { *out = new(HTTPBodyMutation) (*in).DeepCopyInto(*out) } - if in.ContextCache != nil { - in, out := &in.ContextCache, &out.ContextCache - *out = new(ContextCacheSpec) - **out = **in - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AIServiceBackendSpec. @@ -711,21 +706,6 @@ func (in *BackendSecurityPolicyStatus) DeepCopy() *BackendSecurityPolicyStatus { return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *ContextCacheSpec) DeepCopyInto(out *ContextCacheSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ContextCacheSpec. -func (in *ContextCacheSpec) DeepCopy() *ContextCacheSpec { - if in == nil { - return nil - } - out := new(ContextCacheSpec) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *GCPCredentialsFile) DeepCopyInto(out *GCPCredentialsFile) { *out = *in diff --git a/api/v1beta1/ai_service_backend.go b/api/v1beta1/ai_service_backend.go index f45ff7bb52..ed0d3b8573 100644 --- a/api/v1beta1/ai_service_backend.go +++ b/api/v1beta1/ai_service_backend.go @@ -91,42 +91,6 @@ type AIServiceBackendSpec struct { // +kubebuilder:validation:MaxItems=16 HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` - // ContextCache configures context caching for this backend. - // When set, the ai-gateway will automatically resolve or create a context cache - // entry for requests that carry Anthropic-style cache_control markers, and inject the - // resolved cache resource name into the API request. - // Currently only supported for GCP Vertex AI backends. - // +optional - ContextCache *ContextCacheSpec `json:"contextCache,omitempty"` - // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost. // That may be useful for the backend that has a different cost calculation logic. } - -// ContextCacheSpec configures context caching for a backend. -// -// Setting this field enables context caching: cache_control markers in requests are -// resolved against the shared store at URL, and a cache entry is created when absent. -// Leaving it unset disables caching entirely, and markers are ignored. -type ContextCacheSpec struct { - // DefaultTTL is the default time-to-live for newly created cache entries, expressed - // as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. - // Must be at least "60s" per GCP API requirements. - // - // +optional - // +kubebuilder:validation:Pattern=`^[1-9][0-9]*s$` - DefaultTTL string `json:"defaultTTL,omitempty"` - - // URL locates the shared Redis store that records which cachedContents entry a given - // request prefix resolved to. - // - // It is given either as a bare "host:port" - // (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. - // - // If Redis is unreachable at request time, caching is skipped and the request is - // served uncached rather than failed. - // - // +kubebuilder:validation:Required - // +kubebuilder:validation:MinLength=1 - URL string `json:"url"` -} diff --git a/api/v1beta1/zz_generated.deepcopy.go b/api/v1beta1/zz_generated.deepcopy.go index ef2f439333..8a8c91f456 100644 --- a/api/v1beta1/zz_generated.deepcopy.go +++ b/api/v1beta1/zz_generated.deepcopy.go @@ -345,11 +345,6 @@ func (in *AIServiceBackendSpec) DeepCopyInto(out *AIServiceBackendSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ContextCache != nil { - in, out := &in.ContextCache, &out.ContextCache - *out = new(ContextCacheSpec) - **out = **in - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AIServiceBackendSpec. @@ -748,21 +743,6 @@ func (in *BackendSecurityPolicyStatus) DeepCopy() *BackendSecurityPolicyStatus { return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *ContextCacheSpec) DeepCopyInto(out *ContextCacheSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ContextCacheSpec. -func (in *ContextCacheSpec) DeepCopy() *ContextCacheSpec { - if in == nil { - return nil - } - out := new(ContextCacheSpec) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *CredentialOverrideFromDynamicMetadata) DeepCopyInto(out *CredentialOverrideFromDynamicMetadata) { *out = *in 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 317efb48d1..e6ec9a52e7 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...) } @@ -325,6 +340,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( 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/internal/contextcache/gcpcache/gcpcache.go b/internal/contextcache/gcpcache/gcpcache.go index a4abcd8372..3184608e66 100644 --- a/internal/contextcache/gcpcache/gcpcache.go +++ b/internal/contextcache/gcpcache/gcpcache.go @@ -24,6 +24,7 @@ import ( "io" "log/slog" "net/http" + "strconv" "sync" "time" @@ -38,8 +39,9 @@ import ( ) const ( - // defaultTTL is the default cache TTL when none is specified in the cache_control marker. - defaultTTL = "300s" + // 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" @@ -59,6 +61,10 @@ type Resolver struct { // 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 @@ -138,7 +144,7 @@ func (r *Resolver) Close() error { // 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, logger *slog.Logger) *Resolver { +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} } @@ -148,7 +154,13 @@ func New(httpClient *http.Client, store contextcache.Store, auth filterapi.GCPAu if logger == nil { logger = slog.New(slog.DiscardHandler) } - return &Resolver{httpClient: httpClient, store: store, logger: logger, auth: auth} + if defaultTTL <= 0 { + defaultTTL = DefaultTTL + } + return &Resolver{ + httpClient: httpClient, store: store, logger: logger, auth: auth, + defaultTTL: gcpDuration(defaultTTL), + } } // Resolve implements contextcache.CacheResolver. @@ -170,7 +182,7 @@ func (r *Resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion remainderMessages := openAIReq.Messages[breakpoint+1:] // Extract the TTL from the breakpoint message. If not found, fall back to default. - ttl := extractTTL(openAIReq.Messages[breakpoint]) + ttl := extractTTL(openAIReq.Messages[breakpoint], r.defaultTTL) // Build the Gemini cached prefix (tools + system + contents). contents, systemInstruction, err := translator.OpenAIMessagesToGeminiContents(cachedMessages, openAIReq.Model) @@ -358,11 +370,11 @@ func isCacheControlSet(fields *openai.AnthropicContentFields) bool { // 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) string { +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)) + ttl = anthropicTTLToGCP(string(f.CacheControl.TTL), defaultTTL) } } if msg.OfTool != nil { @@ -399,7 +411,7 @@ func extractTTL(msg openai.ChatCompletionMessageParamUnion) string { // anthropicTTLToGCP converts Anthropic cache TTL values to GCP seconds format. // Anthropic supports "5m" and "1h"; unknown values fall back to defaultTTL. -func anthropicTTLToGCP(ttl string) string { +func anthropicTTLToGCP(ttl, defaultTTL string) string { switch ttl { case "5m": return "300s" @@ -414,6 +426,11 @@ func anthropicTTLToGCP(ttl string) string { } } +// 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) diff --git a/internal/contextcache/gcpcache/gcpcache_test.go b/internal/contextcache/gcpcache/gcpcache_test.go index 6500eeaeb0..30c9f552af 100644 --- a/internal/contextcache/gcpcache/gcpcache_test.go +++ b/internal/contextcache/gcpcache/gcpcache_test.go @@ -7,6 +7,8 @@ package gcpcache import ( "context" + "fmt" + "io" "net/http" "net/http/httptest" "sync" @@ -156,7 +158,7 @@ func resolverWithServer(srvURL string, store ...contextcache.Store) *Resolver { return New(&http.Client{ Transport: &redirectTransport{fakeHost: host}, Timeout: 5 * time.Second, - }, s, syncAuth, nil) + }, s, syncAuth, 0, nil) } // memStore is an in-process contextcache.Store for tests. @@ -326,7 +328,7 @@ func TestExtractTTL(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - require.Equal(t, tc.wantTTL, extractTTL(tc.msg)) + require.Equal(t, tc.wantTTL, extractTTL(tc.msg, "300s")) }) } } @@ -620,7 +622,7 @@ func TestResolver_TTLDefault_SentToGoogle(t *testing.T) { var body gcp.CreateCachedContent require.NoError(t, json.Unmarshal(capturedBody, &body)) - assert.Equal(t, defaultTTL, body.TTL) + 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) @@ -712,3 +714,22 @@ func TestResolver_DuplicateCreateRace_KeepsOwnCreate(t *testing.T) { 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) +} diff --git a/internal/contextcache/gcpcache/syncer_test.go b/internal/contextcache/gcpcache/syncer_test.go index ac0fe22877..c415675c71 100644 --- a/internal/contextcache/gcpcache/syncer_test.go +++ b/internal/contextcache/gcpcache/syncer_test.go @@ -267,7 +267,7 @@ func TestResolver_CloseIsIdempotentAndBlocksLaterStart(t *testing.T) { } func TestResolver_StartWithNoopStoreClosesCleanly(t *testing.T) { - r := New(nil, nil, syncAuth, nil) + r := New(nil, nil, syncAuth, 0, nil) r.Start() require.NoError(t, r.Close()) select { diff --git a/internal/controller/controller.go b/internal/controller/controller.go index c7e490b1c3..e5a8061235 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 3b54910af1..37dcee6a63 100644 --- a/internal/controller/gateway.go +++ b/internal/controller/gateway.go @@ -554,13 +554,6 @@ func (c *GatewayController) reconcileFilterConfigSecret( b.Schema = schemaToFilterAPI(backendObj.Spec.APISchema) b.HeaderValueFilters = headerValueFiltersToFilterAPI(backendObj.Spec.HeaderValueFilters) - // Wire context caching configuration when present on the backend object. - if cc := backendObj.Spec.ContextCache; cc != nil { - b.ContextCache = &filterapi.ContextCache{ - URL: cc.URL, - DefaultTTL: cc.DefaultTTL, - } - } } if bsp != nil { diff --git a/internal/controller/gateway_test.go b/internal/controller/gateway_test.go index b8059a9b80..9dd05bf27e 100644 --- a/internal/controller/gateway_test.go +++ b/internal/controller/gateway_test.go @@ -4624,120 +4624,3 @@ func TestGatewayController_warnUndeclaredMetadataNamespaces(t *testing.T) { c.warnUndeclaredMetadataNamespaces(ec, []string{"declared.ns", "missing.ns", "other-missing.ns"}, "gw", "ns") require.Empty(t, logged) } - -// TestGatewayController_reconcileFilterConfigSecret_ContextCache verifies that the -// ContextCache field on an AIServiceBackend is propagated to filterapi.Backend in the -// rendered filter configuration secret. -func TestGatewayController_reconcileFilterConfigSecret_ContextCache(t *testing.T) { - const gwNamespace, someNamespace = "ns", "some-namespace" - fakeClient := requireNewFakeClientWithIndexes(t) - kube := fake2.NewClientset() - ctrl.SetLogger(zap.New(zap.UseFlagOptions(&zap.Options{Development: true, Level: zapcore.DebugLevel}))) - c := newTestGatewayController(fakeClient, kube, ctrl.Log, "envoy-gateway-system", - "docker.io/envoyproxy/ai-gateway-extproc:latest", "info", false, nil, true) - - routes := []aigv1b1.AIGatewayRoute{ - { - ObjectMeta: metav1.ObjectMeta{Name: "gcproute", Namespace: gwNamespace}, - Spec: aigv1b1.AIGatewayRouteSpec{ - Rules: []aigv1b1.AIGatewayRouteRule{ - { - BackendRefs: []aigv1b1.AIGatewayRouteRuleBackendRef{{Name: "gcp-backend"}}, - Matches: []aigv1b1.AIGatewayRouteRuleMatch{{}}, - }, - }, - }, - }, - } - - // Create the AIServiceBackend with ContextCache configured. - err := fakeClient.Create(t.Context(), &aigv1b1.AIServiceBackend{ - ObjectMeta: metav1.ObjectMeta{Name: "gcp-backend", Namespace: gwNamespace}, - Spec: aigv1b1.AIServiceBackendSpec{ - APISchema: aigv1b1.VersionedAPISchema{Name: aigv1b1.APISchemaGCPVertexAI}, - BackendRef: gwapiv1.BackendObjectReference{ - Name: "gcp-service", - Kind: ptr.To(gwapiv1.Kind("Backend")), - Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), - }, - ContextCache: &aigv1b1.ContextCacheSpec{ - DefaultTTL: "600s", - URL: "redis.default.svc.cluster.local:6379", - }, - }, - }) - require.NoError(t, err) - - _, err = c.reconcileFilterConfigSecret(t.Context(), "gw", gwNamespace, someNamespace, routes, nil, "uuid1", nil, nil) - require.NoError(t, err) - - fc := requireFilterConfigFromBundle(t, kube, someNamespace, "gw", gwNamespace) - - // Find the backend that contains "gcp-backend" in its name. - var gcpBackend *filterapi.Backend - for i := range fc.Backends { - if strings.Contains(fc.Backends[i].Name, "gcp-backend") { - gcpBackend = &fc.Backends[i] - break - } - } - require.NotNil(t, gcpBackend, "expected to find gcp-backend in filter config") - require.NotNil(t, gcpBackend.ContextCache, "ContextCache must be populated in filter config") - require.Equal(t, "600s", gcpBackend.ContextCache.DefaultTTL) - require.Equal(t, "redis.default.svc.cluster.local:6379", gcpBackend.ContextCache.URL, - "the redis store location must reach the data plane") -} - -// TestGatewayController_reconcileFilterConfigSecret_ContextCache_NilWhenAbsent verifies that -// when ContextCache is not set on the AIServiceBackend, filterapi.Backend.ContextCache is nil. -func TestGatewayController_reconcileFilterConfigSecret_ContextCache_NilWhenAbsent(t *testing.T) { - const gwNamespace, someNamespace = "ns2", "some-namespace2" - fakeClient := requireNewFakeClientWithIndexes(t) - kube := fake2.NewClientset() - ctrl.SetLogger(zap.New(zap.UseFlagOptions(&zap.Options{Development: true, Level: zapcore.DebugLevel}))) - c := newTestGatewayController(fakeClient, kube, ctrl.Log, "envoy-gateway-system", - "docker.io/envoyproxy/ai-gateway-extproc:latest", "info", false, nil, true) - - routes := []aigv1b1.AIGatewayRoute{ - { - ObjectMeta: metav1.ObjectMeta{Name: "gcproute2", Namespace: gwNamespace}, - Spec: aigv1b1.AIGatewayRouteSpec{ - Rules: []aigv1b1.AIGatewayRouteRule{ - { - BackendRefs: []aigv1b1.AIGatewayRouteRuleBackendRef{{Name: "gcp-backend2"}}, - Matches: []aigv1b1.AIGatewayRouteRuleMatch{{}}, - }, - }, - }, - }, - } - - err := fakeClient.Create(t.Context(), &aigv1b1.AIServiceBackend{ - ObjectMeta: metav1.ObjectMeta{Name: "gcp-backend2", Namespace: gwNamespace}, - Spec: aigv1b1.AIServiceBackendSpec{ - APISchema: aigv1b1.VersionedAPISchema{Name: aigv1b1.APISchemaGCPVertexAI}, - BackendRef: gwapiv1.BackendObjectReference{ - Name: "gcp-service", - Kind: ptr.To(gwapiv1.Kind("Backend")), - Group: ptr.To(gwapiv1.Group("gateway.envoyproxy.io")), - }, - // ContextCache intentionally omitted. - }, - }) - require.NoError(t, err) - - _, err = c.reconcileFilterConfigSecret(t.Context(), "gw", gwNamespace, someNamespace, routes, nil, "uuid2", nil, nil) - require.NoError(t, err) - - fc := requireFilterConfigFromBundle(t, kube, someNamespace, "gw", gwNamespace) - - var gcpBackend *filterapi.Backend - for i := range fc.Backends { - if strings.Contains(fc.Backends[i].Name, "gcp-backend2") { - gcpBackend = &fc.Backends[i] - break - } - } - require.NotNil(t, gcpBackend, "expected to find gcp-backend2 in filter config") - require.Nil(t, gcpBackend.ContextCache, "ContextCache must be nil when not set on the CRD") -} diff --git a/internal/extproc/server.go b/internal/extproc/server.go index 1f1d2fc441..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,8 +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/contextcache/redis" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/internalapi" "github.com/envoyproxy/ai-gateway/internal/redaction" @@ -63,11 +64,14 @@ type Server struct { routerProcessorsPerReqID map[string]Processor routerProcessorsPerReqIDMutex sync.RWMutex uuidFn func() string - // cacheResolvers holds one context-cache resolver per GCP backend, keyed by - // "|". Resolvers are carried across config reloads so a - // filterapi update does not tear down and rebuild the Redis connection pool. The - // Redis URL is part of the key so that repointing a backend at a different Redis - // builds a new resolver rather than reusing a pool aimed at the old one. + // 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 @@ -88,6 +92,13 @@ func NewServer(logger *slog.Logger, enableRedaction bool) (*Server, error) { 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) @@ -102,55 +113,40 @@ func (s *Server) LoadConfig(ctx context.Context, config *filterapi.Config) error logger = slog.New(slog.DiscardHandler) } - // Attach a context-cache resolver to each backend with GCP credentials that - // configures context caching. Resolvers are keyed by backend name and Redis URL and - // reused across reloads, so a config update does not rebuild their Redis connection - // pools; a reused resolver is given the reload's auth handler, since the controller - // rotates the token. Resolvers whose backend disappeared, dropped its cache config, or - // moved to another Redis are closed, which stops their background syncer. + // 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 _, ok := rb.Handler.(filterapi.GCPAuthHandler); !ok { + if s.cacheStore == nil || rb.Backend.Schema.Name != filterapi.APISchemaGCPVertexAI { continue } - cc := rb.Backend.ContextCache - if cc == nil || cc.URL == "" { + auth, ok := rb.Handler.(filterapi.GCPAuthHandler) + if !ok { continue } - key := rb.Backend.Name + "|" + cc.URL - resolver, ok := s.cacheResolvers[key] - if !ok { - store, err := redis.NewStore(cc.URL) - if err != nil { - // A bad URL disables caching for this backend rather than failing the - // reload: the store is an optimization, and rejecting the whole config - // would take down backends that have nothing to do with caching. - logger.Warn("invalid context cache redis url, serving backend uncached", - slog.String("backend", rb.Backend.Name), slog.String("error", err.Error())) - continue - } - resolver = gcpcache.New(nil, store, rb.Handler.(filterapi.GCPAuthHandler), logger) + name := rb.Backend.Name + resolver, ok := s.cacheResolvers[name] + if ok { + resolver.SetAuth(auth) } else { - // The handler is rebuilt on every reload (the controller rotates the access - // token), so hand the reused resolver the current one. - resolver.SetAuth(rb.Handler.(filterapi.GCPAuthHandler)) + resolver = gcpcache.New(nil, s.cacheStore, auth, s.cacheDefaultTTL, logger) } - // Start is a no-op for a resolver already running, so reused resolvers are - // unaffected. It runs under the resolver's own context, not ctx, which is - // cancelled when this reload finishes. + // 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[key] = resolver + live[name] = resolver rb.CacheResolver = resolver } - // Stop the background work of resolvers that were not carried into this config: their - // backend was removed, lost its cache config, or was repointed at another Redis. - for key, resolver := range s.cacheResolvers { - if _, ok := live[key]; ok { - continue + for name, resolver := range s.cacheResolvers { + if _, ok := live[name]; !ok { + _ = resolver.Close() } - _ = resolver.Close() } s.cacheResolvers = live diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go index bd97a37a69..f0d17c6bd0 100644 --- a/internal/extproc/server_test.go +++ b/internal/extproc/server_test.go @@ -25,6 +25,7 @@ import ( "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" @@ -55,186 +56,91 @@ func TestServer_LoadConfig(t *testing.T) { require.NotNil(t, s.config) } -// TestServer_LoadConfig_GCPCacheResolver verifies that a CacheResolver is attached -// to a GCP backend only when ContextCache is configured with a store URL. -func TestServer_LoadConfig_GCPCacheResolver(t *testing.T) { +// 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) - - config := &filterapi.Config{ - Backends: []filterapi.Backend{ - { - Name: "gcp-caching-on", - Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, - Auth: &filterapi.BackendAuth{ - GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, - }, - ContextCache: &filterapi.ContextCache{URL: "localhost:6379", DefaultTTL: "600s"}, - }, - { - Name: "gcp-caching-no-url", - Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, - Auth: &filterapi.BackendAuth{ - GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, - }, - ContextCache: &filterapi.ContextCache{}, - }, - { - Name: "gcp-no-caching-field", - Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, - Auth: &filterapi.BackendAuth{ - GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}, - }, - // ContextCache is nil. - }, - { - Name: "non-gcp", - Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}, - // No auth handler → not a GCPAuthHandler. - }, - }, - } - - err := s.LoadConfig(t.Context(), config) - require.NoError(t, err) - - rb := s.config.Backends["gcp-caching-on"] - require.NotNil(t, rb.CacheResolver, "CacheResolver must be set when ContextCache has a store URL") - - rb = s.config.Backends["gcp-caching-no-url"] - require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when ContextCache has no store URL") - - rb = s.config.Backends["gcp-no-caching-field"] - require.Nil(t, rb.CacheResolver, "CacheResolver must be nil when ContextCache is absent") - - rb = s.config.Backends["non-gcp"] - require.Nil(t, rb.CacheResolver, "CacheResolver must be nil for non-GCP backends") + return s } -// gcpBackendWithCaching builds a GCP backend whose context caching points at redisURL. -func gcpBackendWithCaching(name, redisURL string) filterapi.Backend { +// 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"}}, - ContextCache: &filterapi.ContextCache{URL: redisURL, DefaultTTL: "600s"}, + Name: name, + Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaGCPVertexAI}, + Auth: &filterapi.BackendAuth{GCPAuth: &filterapi.GCPAuth{AccessToken: "token"}}, } } -// A resolver owns a Redis connection pool. Rebuilding it on every filterapi update would -// discard live connections on a hot path that fires whenever any unrelated backend changes. -func TestServer_LoadConfig_GCPCacheResolver_ReusedAcrossReloads(t *testing.T) { - s := &Server{} - t.Cleanup(s.Close) - config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}} - - require.NoError(t, s.LoadConfig(t.Context(), config)) - first := s.config.Backends["gcp"].CacheResolver - require.NotNil(t, first) - - require.NoError(t, s.LoadConfig(t.Context(), config)) - require.Same(t, first, s.config.Backends["gcp"].CacheResolver, - "an unchanged backend must keep its resolver, and with it its connection pool") -} - -// A reused resolver must pick up the handler rebuilt on reload, since the controller -// rotates the access token without changing the Redis URL. -func TestServer_LoadConfig_GCPCacheResolver_ReusedResolverGetsNewHandler(t *testing.T) { - s := &Server{} - t.Cleanup(s.Close) - config := &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6379")}} - - require.NoError(t, s.LoadConfig(t.Context(), config)) - first := s.config.Backends["gcp"].Handler +// 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.NoError(t, s.LoadConfig(t.Context(), config)) - second := s.config.Backends["gcp"].Handler - require.NotSame(t, first, second, "the handler is rebuilt on every reload") - require.Equal(t, second, s.cacheResolvers["gcp|localhost:6379"].GetAuth(), - "the reused resolver must hold the current handler") + 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) } -// Repointing a backend at a different Redis must not keep reusing a pool dialing the old one. -func TestServer_LoadConfig_GCPCacheResolver_RebuiltWhenRedisURLChanges(t *testing.T) { +// 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{gcpBackendWithCaching("gcp", "localhost:6379")}})) - first := s.config.Backends["gcp"].CacheResolver - require.NotNil(t, first) - - require.NoError(t, s.LoadConfig(t.Context(), - &filterapi.Config{Backends: []filterapi.Backend{gcpBackendWithCaching("gcp", "localhost:6380")}})) - require.NotSame(t, first, s.config.Backends["gcp"].CacheResolver, - "a new Redis URL must produce a new resolver") + 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) } -// Resolvers for backends that vanish must not accumulate: the map is rebuilt each reload. -func TestServer_LoadConfig_GCPCacheResolver_DroppedWhenBackendDisappears(t *testing.T) { - s := &Server{} - t.Cleanup(s.Close) +// 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(), &filterapi.Config{Backends: []filterapi.Backend{ - gcpBackendWithCaching("gcp-a", "localhost:6379"), - gcpBackendWithCaching("gcp-b", "localhost:6379"), - }})) - require.Len(t, s.cacheResolvers, 2) - - require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ - gcpBackendWithCaching("gcp-a", "localhost:6379"), - }})) - require.Len(t, s.cacheResolvers, 1, "the departed backend's resolver must be released") -} + require.NoError(t, s.LoadConfig(t.Context(), config)) + first := s.config.Backends["gemini"].CacheResolver + firstHandler := s.config.Backends["gemini"].Handler -// A malformed Redis URL disables caching for that backend rather than failing the whole -// reload, which would take down backends that have nothing to do with caching. -func TestServer_LoadConfig_GCPCacheResolver_BadRedisURLDoesNotFailReload(t *testing.T) { - s := &Server{} - t.Cleanup(s.Close) - config := &filterapi.Config{Backends: []filterapi.Backend{ - gcpBackendWithCaching("gcp-bad-redis", "http://not-a-redis-url"), - gcpBackendWithCaching("gcp-ok", "localhost:6379"), - }} - - require.NoError(t, s.LoadConfig(t.Context(), config), "a bad redis url must not fail the reload") - require.Nil(t, s.config.Backends["gcp-bad-redis"].CacheResolver, - "a backend whose store cannot be built is served uncached") - require.NotNil(t, s.config.Backends["gcp-ok"].CacheResolver, - "an unrelated backend must keep its resolver") + 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, or every -// reload that drops or repoints a backend would leak a goroutine. -func TestServer_LoadConfig_GCPCacheResolver_EvictedResolverIsClosed(t *testing.T) { +// 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{ - gcpBackendWithCaching("gcp-a", "localhost:6379"), - gcpBackendWithCaching("gcp-b", "localhost:6379"), - }})) - kept := s.cacheResolvers["gcp-a|localhost:6379"] - - require.NoError(t, s.LoadConfig(t.Context(), &filterapi.Config{Backends: []filterapi.Backend{ - gcpBackendWithCaching("gcp-a", "localhost:6379"), + gcpBackend("a"), gcpBackend("b"), }})) - // Closing the one resolver still in the config must leave nothing running. If the - // evicted resolver had not been closed by the reload, its syncer would remain. + 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 not be tied -// to it, or they would stop seconds after every config load. +// 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{ - gcpBackendWithCaching("gcp", "localhost:6379"), - }})) + 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") diff --git a/internal/filterapi/filterconfig.go b/internal/filterapi/filterconfig.go index f82d151458..f2b7775b6f 100644 --- a/internal/filterapi/filterconfig.go +++ b/internal/filterapi/filterconfig.go @@ -207,27 +207,6 @@ type Backend struct { // HeaderValueFilters filter individual values out of multi-valued request headers before sending // the request to the backend. Optional. HeaderValueFilters []HTTPHeaderValueFilter `json:"headerValueFilters,omitempty"` - // ContextCache configures context caching for this backend. Optional. - // When non-nil, the extproc will resolve or create a GCP cachedContents entry - // for requests that carry cache_control markers. - ContextCache *ContextCache `json:"contextCache,omitempty"` -} - -// ContextCache configures context caching for a GCP Vertex AI backend. -// -// Its presence enables caching: markers are resolved against the store at URL and a -// cache entry is created when absent. -type ContextCache struct { - // DefaultTTL is the default time-to-live for newly created cache entries expressed - // as a duration string (e.g. "600s"). When empty, the resolver default (300s) is used. - // +optional - DefaultTTL string `json:"defaultTTL,omitempty"` - // URL points at the shared Redis store that records which cachedContents entry a - // given prefix resolved to. Once one replica has created a cache, the others reuse - // it. Replicas that miss the same prefix at the same moment may each create one. - // - // It is either a bare "host:port" or a full "redis://" URL. - URL string `json:"url"` } // BackendAuth corresponds partially to BackendSecurityPolicy in api/v1alpha1/api.go. diff --git a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml index 296007ed91..64a3fa17c1 100644 --- a/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml +++ b/manifests/charts/ai-gateway-crds-helm/templates/aigateway.envoyproxy.io_aiservicebackends.yaml @@ -224,36 +224,6 @@ spec: - path x-kubernetes-list-type: map type: object - contextCache: - description: |- - ContextCache configures context caching for this backend. - When set, the ai-gateway will automatically resolve or create a context cache - entry for requests that carry Anthropic-style cache_control markers, and inject the - resolved cache resource name into the API request. - Currently only supported for GCP Vertex AI backends. - properties: - defaultTTL: - description: |- - DefaultTTL is the default time-to-live for newly created cache entries, expressed - as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. - Must be at least "60s" per GCP API requirements. - pattern: ^[1-9][0-9]*s$ - type: string - url: - description: |- - URL locates the shared Redis store that records which cachedContents entry a given - request prefix resolved to. - - It is given either as a bare "host:port" - (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. - - If Redis is unreachable at request time, caching is skipped and the request is - served uncached rather than failed. - minLength: 1 - type: string - required: - - url - type: object headerMutation: description: |- HeaderMutation defines the mutation of HTTP headers that will be applied to the request @@ -671,36 +641,6 @@ spec: - path x-kubernetes-list-type: map type: object - contextCache: - description: |- - ContextCache configures context caching for this backend. - When set, the ai-gateway will automatically resolve or create a context cache - entry for requests that carry Anthropic-style cache_control markers, and inject the - resolved cache resource name into the API request. - Currently only supported for GCP Vertex AI backends. - properties: - defaultTTL: - description: |- - DefaultTTL is the default time-to-live for newly created cache entries, expressed - as a GCP duration string (e.g. "600s"). When empty, the default of 300s is used. - Must be at least "60s" per GCP API requirements. - pattern: ^[1-9][0-9]*s$ - type: string - url: - description: |- - URL locates the shared Redis store that records which cachedContents entry a given - request prefix resolved to. - - It is given either as a bare "host:port" - (e.g. "redis.default.svc.cluster.local:6379") or as a full "redis://" URL. - - If Redis is unreachable at request time, caching is skipped and the request is - served uncached rather than failed. - minLength: 1 - type: string - required: - - url - type: object headerMutation: description: |- HeaderMutation defines the mutation of HTTP headers that will be applied to the request 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: diff --git a/site/docs/api/api.mdx b/site/docs/api/api.mdx index 1139e574f9..03621506b7 100644 --- a/site/docs/api/api.mdx +++ b/site/docs/api/api.mdx @@ -579,7 +579,6 @@ QuotaPolicyList contains a list of QuotaPolicy - [BackendSecurityPolicySpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicyspec) - [BackendSecurityPolicyStatus](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicystatus) - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1alpha1-backendsecuritypolicytype) -- [ContextCacheSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-contextcachespec) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpcredentialsfile) - [GCPOIDCExchangeToken](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpoidcexchangetoken) - [GCPServiceAccountImpersonationConfig](#github-com-envoyproxy-ai-gateway-api-v1alpha1-gcpserviceaccountimpersonationconfig) @@ -857,11 +856,6 @@ AIServiceBackendSpec details the AIServiceBackend configuration. type="[HTTPBodyMutation](#github-com-envoyproxy-ai-gateway-api-v1alpha1-httpbodymutation)" required="false" description="BodyMutation defines the mutation of HTTP request body JSON fields that will be applied to the request
before sending it to the backend." -/> @@ -1387,36 +1381,6 @@ BackendSecurityPolicyType specifies the type of auth mechanism used to access a required="false" description="" /> -#### ContextCacheSpec - - - -**Appears in:** -- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1alpha1-aiservicebackendspec) - -ContextCacheSpec configures context caching for a backend. - -Setting this field enables context caching: cache_control markers in requests are -resolved against the shared store at URL, and a cache entry is created when absent. -Leaving it unset disables caching entirely, and markers are ignored. - -##### Fields - - - - - - #### GCPCredentialsFile @@ -3196,7 +3160,6 @@ MCPRouteList contains a list of MCPRoute. - [BackendSecurityPolicySpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicyspec) - [BackendSecurityPolicyStatus](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicystatus) - [BackendSecurityPolicyType](#github-com-envoyproxy-ai-gateway-api-v1beta1-backendsecuritypolicytype) -- [ContextCacheSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-contextcachespec) - [CredentialOverrideFromDynamicMetadata](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromdynamicmetadata) - [CredentialOverrideFromRequestHeaders](#github-com-envoyproxy-ai-gateway-api-v1beta1-credentialoverridefromrequestheaders) - [GCPCredentialsFile](#github-com-envoyproxy-ai-gateway-api-v1beta1-gcpcredentialsfile) @@ -3478,11 +3441,6 @@ AIServiceBackendSpec details the AIServiceBackend configuration. type="[HTTPHeaderValueFilter](#github-com-envoyproxy-ai-gateway-api-v1beta1-httpheadervaluefilter) array" required="false" description="HeaderValueFilters filter individual values out of multi-valued request headers before sending
the request to this backend, so that one value an upstream provider does not accept does not
fail the whole request. At most one filter per header name.
Only GCPAnthropic and AWSAnthropic backends honor this, and today only for the
`anthropic-beta` header — a filter naming any other header is accepted but has no effect.
Setting it on a backend with any other schema is rejected." -/> @@ -4046,36 +4004,6 @@ BackendSecurityPolicyType specifies the type of auth mechanism used to access a required="false" description="" /> -#### ContextCacheSpec - - - -**Appears in:** -- [AIServiceBackendSpec](#github-com-envoyproxy-ai-gateway-api-v1beta1-aiservicebackendspec) - -ContextCacheSpec configures context caching for a backend. - -Setting this field enables context caching: cache_control markers in requests are -resolved against the shared store at URL, and a cache entry is created when absent. -Leaving it unset disables caching entirely, and markers are ignored. - -##### Fields - - - - - - #### CredentialOverrideFromDynamicMetadata diff --git a/tests/crdcel/main_test.go b/tests/crdcel/main_test.go index d9182b9092..6a3511a104 100644 --- a/tests/crdcel/main_test.go +++ b/tests/crdcel/main_test.go @@ -95,17 +95,6 @@ func TestAIServiceBackends(t *testing.T) { }{ {name: "basic.yaml"}, {name: "anthropic-schema.yaml"}, - {name: "context-cache-enabled.yaml"}, - { - name: "context-cache-invalid-ttl.yaml", - expErr: "spec.contextCache.defaultTTL", - }, - { - // Url should not be empty, so a missing - // url reaches the API server as "" and fails the MinLength rule. - name: "context-cache-missing-url.yaml", - expErr: "spec.contextCache.url in body should be at least 1 chars long", - }, {name: "basic-eg-backend-aws.yaml"}, {name: "aws-openai-schema.yaml"}, {name: "basic-eg-backend-azure.yaml"}, diff --git a/tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml deleted file mode 100644 index db43d02160..0000000000 --- a/tests/crdcel/testdata/aiservicebackends/context-cache-enabled.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# 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. - -apiVersion: aigateway.envoyproxy.io/v1beta1 -kind: AIServiceBackend -metadata: - name: gcp-cache-backend - namespace: default -spec: - schema: - name: GCPVertexAI - backendRef: - name: gcp-service - kind: Backend - group: gateway.envoyproxy.io - port: 443 - contextCache: - url: "redis.default.svc.cluster.local:6379" - defaultTTL: "600s" diff --git a/tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml deleted file mode 100644 index 4744640280..0000000000 --- a/tests/crdcel/testdata/aiservicebackends/context-cache-invalid-ttl.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# 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. - -apiVersion: aigateway.envoyproxy.io/v1beta1 -kind: AIServiceBackend -metadata: - name: gcp-cache-backend-bad-ttl - namespace: default -spec: - schema: - name: GCPVertexAI - backendRef: - name: gcp-service - kind: Backend - group: gateway.envoyproxy.io - port: 443 - contextCache: - url: "redis.default.svc.cluster.local:6379" - defaultTTL: "not-a-duration" diff --git a/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml b/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml deleted file mode 100644 index 26edc20522..0000000000 --- a/tests/crdcel/testdata/aiservicebackends/context-cache-missing-url.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# 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. - -apiVersion: aigateway.envoyproxy.io/v1beta1 -kind: AIServiceBackend -metadata: - name: gcp-cache-backend-no-url - namespace: default -spec: - schema: - name: GCPVertexAI - backendRef: - name: gcp-service - kind: Backend - group: gateway.envoyproxy.io - port: 443 - contextCache: - defaultTTL: "600s" From b50eeecf60fd7d453acb7601c39b3260fe3096a1 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Wed, 7 Oct 2026 23:01:15 +0530 Subject: [PATCH 14/17] translator/gcpvertexai: remove the cachedContent vendor field Users do not set a Vertex AI cache name themselves; the gateway's context-cache resolver is the only thing that does. Remove the cachedContent field from GCPVertexAIVendorFields along with the code that supported it: - the RequestBody check rejecting cachedContent combined with cache_control markers, and its gcpRequestHasCacheControlMarkers helper, which had no other caller; - copying the vendor field into the Gemini request, and the condition that let it override the resolved cache name; - the processor guard that skipped resolution when a request named its own cache. The field was added on this branch and is not in a release. A client that still sends it has it ignored. Signed-off-by: Sivanantham Chinnaiyan --- internal/apischema/openai/openai.go | 9 - internal/extproc/processor_impl.go | 4 +- internal/extproc/processor_impl_test.go | 22 --- internal/translator/openai_gcpvertexai.go | 55 +----- .../translator/openai_gcpvertexai_test.go | 179 ------------------ 5 files changed, 2 insertions(+), 267 deletions(-) diff --git a/internal/apischema/openai/openai.go b/internal/apischema/openai/openai.go index cbb65aac03..b992240e02 100644 --- a/internal/apischema/openai/openai.go +++ b/internal/apischema/openai/openai.go @@ -2069,15 +2069,6 @@ type GCPVertexAIVendorFields struct { // // https://cloud.google.com/vertex-ai/docs/reference/rest/v1/SafetySetting SafetySettings []*genai.SafetySetting `json:"safetySettings,omitzero"` - - // CachedContent specifies a pre-existing Vertex AI cached content resource to use as context. - // Format: "projects/{project}/locations/{location}/cachedContents/{cache_id}" - // - // When provided, cache_control markers on messages must be absent. The cache reference is - // sent directly to Gemini without calling any cache resolution logic. - // - // https://cloud.google.com/vertex-ai/docs/context-cache/context-cache-overview - CachedContent string `json:"cachedContent,omitzero"` } // GCPVertexAIGenerationConfig represents Gemini generation configuration options. diff --git a/internal/extproc/processor_impl.go b/internal/extproc/processor_impl.go index 5ce8341d7f..0d2f478d0f 100644 --- a/internal/extproc/processor_impl.go +++ b/internal/extproc/processor_impl.go @@ -381,9 +381,7 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque // the resolved cache name and filtered message list. if u.cacheResolver != nil { if cacheSetter, ok := u.translator.(translator.ContextCacheSetter); ok { - // Skip requests that name their own cache. - if req, ok := any(u.parent.originalRequestBody).(*openai.ChatCompletionRequest); ok && - (req.GCPVertexAIVendorFields == nil || req.CachedContent == "") { + 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())) diff --git a/internal/extproc/processor_impl_test.go b/internal/extproc/processor_impl_test.go index 5729af26ed..dde75cd5c3 100644 --- a/internal/extproc/processor_impl_test.go +++ b/internal/extproc/processor_impl_test.go @@ -2723,28 +2723,6 @@ func Test_upstreamProcessor_CacheResolver_hitIgnoresTokenCount(t *testing.T) { require.Zero(t, tr.setCacheResult.WriteTokenCount) } -// A request that names its own cache must not be resolved, so no cache is created for it. -func Test_upstreamProcessor_CacheResolver_explicitCachedContentSkips(t *testing.T) { - someBody := bodyFromModel(t, "gemini-1.5-pro", false, nil) - var body openai.ChatCompletionRequest - require.NoError(t, json.Unmarshal(someBody, &body)) - body.GCPVertexAIVendorFields = &openai.GCPVertexAIVendorFields{CachedContent: "c/explicit"} - - cr := &mockCacheResolver{} - 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, - } - _, _ = p.ProcessRequestHeaders(t.Context(), nil) - require.False(t, cr.called) - require.Nil(t, tr.setCacheResult) -} - // 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) diff --git a/internal/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index 2be22614c6..677b746335 100644 --- a/internal/translator/openai_gcpvertexai.go +++ b/internal/translator/openai_gcpvertexai.go @@ -113,12 +113,6 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) RequestBody(_ []byte, op o.requestModel = o.modelNameOverride } - // Validate that explicit cachedContent and cache_control markers are not used together. - if openAIReq.GCPVertexAIVendorFields != nil && openAIReq.CachedContent != "" && - gcpRequestHasCacheControlMarkers(openAIReq) { - return nil, nil, fmt.Errorf("%w: cannot specify both cache_control on messages and explicit cachedContent field", internalapi.ErrMalformedRequest) - } - // Set streaming flag. o.stream = openAIReq.Stream @@ -596,9 +590,7 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes o.applyVendorSpecificFields(openAIReq, &gcr, requestModel) // Inject resolved cache name from the context-cache resolver (set via SetContextCacheResult). - // This happens after vendor-field processing so an explicit vendor cachedContent still wins - // (the validation in RequestBody already rejected the mixed case). - if o.cacheResult != nil && gcr.CachedContent == "" { + 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. @@ -610,48 +602,6 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes return &gcr, nil } -// gcpRequestHasCacheControlMarkers returns true if any message content part in the request -// contains a cache_control marker (Anthropic-style ephemeral cache breakpoint). -func gcpRequestHasCacheControlMarkers(openAIReq *openai.ChatCompletionRequest) bool { - for i := range openAIReq.Messages { - msg := &openAIReq.Messages[i] - // Check tool messages (AnthropicContentFields inline on message itself). - if msg.OfTool != nil && msg.OfTool.AnthropicContentFields != nil && - isCacheEnabled(msg.OfTool.AnthropicContentFields) { - return true - } - // Check system messages (content is ContentUnion — []ChatCompletionContentPartTextParam or string). - if msg.OfSystem != nil { - if parts, ok := msg.OfSystem.Content.Value.([]openai.ChatCompletionContentPartTextParam); ok { - for j := range parts { - if isCacheEnabled(parts[j].AnthropicContentFields) { - return true - } - } - } - } - // Check user messages (content is StringOrUserRoleContentUnion — - // []ChatCompletionContentPartUserUnionParam or string). - if msg.OfUser != nil { - if parts, ok := msg.OfUser.Content.Value.([]openai.ChatCompletionContentPartUserUnionParam); ok { - for j := range parts { - part := &parts[j] - if part.OfText != nil && isCacheEnabled(part.OfText.AnthropicContentFields) { - return true - } - if part.OfImageURL != nil && isCacheEnabled(part.OfImageURL.AnthropicContentFields) { - return true - } - if part.OfInputAudio != nil && isCacheEnabled(part.OfInputAudio.AnthropicContentFields) { - return true - } - } - } - } - } - return false -} - // applyVendorSpecificFields applies GCP Vertex AI vendor-specific fields to the Gemini request. // These fields allow users to access advanced GCP-specific features not available in the OpenAI API. // Vendor fields override any conflicting fields that were set during the standard translation process. @@ -674,9 +624,6 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) applyVendorSpecificField if gcpVendorFields.SafetySettings != nil { gcr.SafetySettings = gcpVendorFields.SafetySettings } - if gcpVendorFields.CachedContent != "" { - gcr.CachedContent = gcpVendorFields.CachedContent - } } // SetContextCacheResult implements [ContextCacheSetter]. The upstream processor calls this diff --git a/internal/translator/openai_gcpvertexai_test.go b/internal/translator/openai_gcpvertexai_test.go index 60896bbbf4..75a7091584 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -12,8 +12,6 @@ import ( "strings" "testing" - anthropic "github.com/anthropics/anthropic-sdk-go" - "github.com/anthropics/anthropic-sdk-go/shared/constant" corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" "github.com/google/go-cmp/cmp" @@ -782,71 +780,6 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) }, wantBody: wantBdyWithEnterpriseWebSearch, }, - { - name: "Request with explicit cachedContent vendor field", - input: openai.ChatCompletionRequest{ - Model: "gemini-1.5-pro", - MaxTokens: ptr.To(int64(1024)), - Messages: []openai.ChatCompletionMessageParamUnion{ - { - OfUser: &openai.ChatCompletionUserMessageParam{ - Role: openai.ChatMessageRoleUser, - Content: openai.StringOrUserRoleContentUnion{Value: "Summarize the document."}, - }, - }, - }, - GCPVertexAIVendorFields: &openai.GCPVertexAIVendorFields{ - CachedContent: "projects/my-project/locations/us-central1/cachedContents/abc123", - }, - }, - onRetry: false, - wantError: false, - wantHeaderMut: []internalapi.Header{ - {":path", "publishers/google/models/gemini-1.5-pro:generateContent"}, - }, - wantBody: []byte(`{ - "contents": [{"parts": [{"text": "Summarize the document."}], "role": "user"}], - "tools": null, - "generationConfig": {"maxOutputTokens": 1024}, - "cachedContent": "projects/my-project/locations/us-central1/cachedContents/abc123" - }`), - }, - { - name: "Request with both cache_control markers and explicit cachedContent returns error", - input: openai.ChatCompletionRequest{ - Model: "gemini-1.5-pro", - MaxTokens: ptr.To(int64(1024)), - Messages: []openai.ChatCompletionMessageParamUnion{ - { - OfSystem: &openai.ChatCompletionSystemMessageParam{ - Role: openai.ChatMessageRoleSystem, - Content: openai.ContentUnion{Value: []openai.ChatCompletionContentPartTextParam{ - { - Type: "text", - Text: "You are a helpful assistant.", - AnthropicContentFields: &openai.AnthropicContentFields{ - CacheControl: anthropic.CacheControlEphemeralParam{ - Type: constant.ValueOf[constant.Ephemeral](), - }, - }, - }, - }}, - }, - }, - { - OfUser: &openai.ChatCompletionUserMessageParam{ - Role: openai.ChatMessageRoleUser, - Content: openai.StringOrUserRoleContentUnion{Value: "Hello"}, - }, - }, - }, - GCPVertexAIVendorFields: &openai.GCPVertexAIVendorFields{ - CachedContent: "projects/my-project/locations/us-central1/cachedContents/abc123", - }, - }, - onRetry: true, - wantError: true, - }, } for _, tc := range tests { @@ -886,118 +819,6 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) } } -func TestGCPRequestHasCacheControlMarkers(t *testing.T) { - ephemeral := anthropic.CacheControlEphemeralParam{Type: constant.ValueOf[constant.Ephemeral]()} - - tests := []struct { - name string - req openai.ChatCompletionRequest - want bool - }{ - { - name: "no cache_control markers", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfUser: &openai.ChatCompletionUserMessageParam{ - Role: openai.ChatMessageRoleUser, - Content: openai.StringOrUserRoleContentUnion{Value: "hello"}, - }}, - }, - }, - want: false, - }, - { - name: "cache_control on system message content part", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfSystem: &openai.ChatCompletionSystemMessageParam{ - Role: openai.ChatMessageRoleSystem, - Content: openai.ContentUnion{Value: []openai.ChatCompletionContentPartTextParam{ - { - Type: "text", - Text: "You are helpful.", - AnthropicContentFields: &openai.AnthropicContentFields{ - CacheControl: ephemeral, - }, - }, - }}, - }}, - }, - }, - want: true, - }, - { - name: "cache_control on user message text content part", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfUser: &openai.ChatCompletionUserMessageParam{ - Role: openai.ChatMessageRoleUser, - Content: openai.StringOrUserRoleContentUnion{ - Value: []openai.ChatCompletionContentPartUserUnionParam{ - {OfText: &openai.ChatCompletionContentPartTextParam{ - Type: "text", - Text: "large context", - AnthropicContentFields: &openai.AnthropicContentFields{ - CacheControl: ephemeral, - }, - }}, - }, - }, - }}, - }, - }, - want: true, - }, - { - name: "cache_control on tool message", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfTool: &openai.ChatCompletionToolMessageParam{ - Role: openai.ChatMessageRoleTool, - ToolCallID: "call_1", - Content: openai.ContentUnion{Value: "result"}, - AnthropicContentFields: &openai.AnthropicContentFields{ - CacheControl: ephemeral, - }, - }}, - }, - }, - want: true, - }, - { - name: "string user message has no markers", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfUser: &openai.ChatCompletionUserMessageParam{ - Role: openai.ChatMessageRoleUser, - Content: openai.StringOrUserRoleContentUnion{Value: "plain string"}, - }}, - }, - }, - want: false, - }, - { - name: "system string content has no markers", - req: openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessageParamUnion{ - {OfSystem: &openai.ChatCompletionSystemMessageParam{ - Role: openai.ChatMessageRoleSystem, - Content: openai.ContentUnion{Value: "plain system string"}, - }}, - }, - }, - want: false, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got := gcpRequestHasCacheControlMarkers(&tc.req) - require.Equal(t, tc.want, got) - }) - } -} - func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_ResponseHeaders(t *testing.T) { tests := []struct { name string From 89501938beb4734f987c0cfee96f2a52a08e817e Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Thu, 8 Oct 2026 15:06:26 +0530 Subject: [PATCH 15/17] gcpcache: sort map keys when computing the cache key The cache key is a SHA-256 over the JSON of the cached prefix. It was encoded with internal/json.Marshal, which uses sonic without SortMapKeys, so map keys came out in Go's random map iteration order. Tool schema properties and function call args and results are maps, so any prefix containing them could get a different key on each request. Each new key missed the store and created another billed cache, which made caching largely ineffective for tool-heavy prompts. Add internal/json.MarshalSorted, a sonic config with SortMapKeys, and use it in computeCacheKey. Add a test that computes the key repeatedly for a prefix with a 50-property tool schema and 50 function call args; it fails with the old encoder. Keys change for prefixes that contain maps, so each such prefix misses once after deploy and creates a new cache. Text-only prefixes keep their keys. Old store entries expire with their caches. Signed-off-by: Sivanantham Chinnaiyan --- internal/contextcache/gcpcache/gcpcache.go | 4 +++- .../contextcache/gcpcache/gcpcache_test.go | 24 +++++++++++++++++++ internal/json/json.go | 10 +++++++- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/internal/contextcache/gcpcache/gcpcache.go b/internal/contextcache/gcpcache/gcpcache.go index 3184608e66..17c6beca13 100644 --- a/internal/contextcache/gcpcache/gcpcache.go +++ b/internal/contextcache/gcpcache/gcpcache.go @@ -468,7 +468,9 @@ func computeCacheKey(project, region, model string, contents []genai.Content, sy SystemInstruction: systemInstruction, Tools: tools, } - b, err := json.Marshal(input) + // 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) } diff --git a/internal/contextcache/gcpcache/gcpcache_test.go b/internal/contextcache/gcpcache/gcpcache_test.go index 30c9f552af..7bb56a0e7d 100644 --- a/internal/contextcache/gcpcache/gcpcache_test.go +++ b/internal/contextcache/gcpcache/gcpcache_test.go @@ -21,6 +21,7 @@ import ( "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" @@ -733,3 +734,26 @@ func TestResolver_ConfiguredDefaultTTL_SentToGoogle(t *testing.T) { 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) + } +} 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. From 72d508cbf28c625406ad4b958dfd67cc2bcceaeb Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Thu, 8 Oct 2026 21:26:51 +0530 Subject: [PATCH 16/17] wip Signed-off-by: Sivanantham Chinnaiyan --- internal/contextcache/gcpcache/gcpcache.go | 11 ++++++++ .../contextcache/gcpcache/gcpcache_test.go | 28 +++++++++++++++++++ internal/extproc/processor_impl.go | 12 ++++++++ 3 files changed, 51 insertions(+) diff --git a/internal/contextcache/gcpcache/gcpcache.go b/internal/contextcache/gcpcache/gcpcache.go index 17c6beca13..8cd8c0cb83 100644 --- a/internal/contextcache/gcpcache/gcpcache.go +++ b/internal/contextcache/gcpcache/gcpcache.go @@ -204,6 +204,10 @@ func (r *Resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // 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, @@ -213,10 +217,17 @@ func (r *Resolver) Resolve(ctx context.Context, openAIReq *openai.ChatCompletion // 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. diff --git a/internal/contextcache/gcpcache/gcpcache_test.go b/internal/contextcache/gcpcache/gcpcache_test.go index 7bb56a0e7d..2b6ee8abd0 100644 --- a/internal/contextcache/gcpcache/gcpcache_test.go +++ b/internal/contextcache/gcpcache/gcpcache_test.go @@ -6,9 +6,11 @@ package gcpcache import ( + "bytes" "context" "fmt" "io" + "log/slog" "net/http" "net/http/httptest" "sync" @@ -757,3 +759,29 @@ func TestComputeCacheKey_StableWithMaps(t *testing.T) { 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/extproc/processor_impl.go b/internal/extproc/processor_impl.go index 0d2f478d0f..eabc436b8e 100644 --- a/internal/extproc/processor_impl.go +++ b/internal/extproc/processor_impl.go @@ -35,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" ) @@ -379,6 +380,8 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque // 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 { @@ -399,12 +402,21 @@ func (u *upstreamProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessReque 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" From de15b5e32befdbcabe545364fdb5fff03016ec69 Mon Sep 17 00:00:00 2001 From: Sivanantham Chinnaiyan Date: Sun, 11 Oct 2026 14:45:29 +0530 Subject: [PATCH 17/17] backendauth: keep the ADC token source off the caller's context newGCPHandler passed its ctx to google.FindDefaultCredentials. The returned token source keeps that context and uses it for every later token refresh, because oauth2.TokenSource.Token takes no context. Config reloads build handlers with a per-tick context that the watcher cancels as soon as the reload returns. The same happens on startup when the first load is retried on a tick. After that, every token fetch failed with "context canceled", on the request path and in the context cache syncer. AWS is unaffected because it passes the request context to Retrieve on each call. Build the token source on context.Background() carrying only the proxy-aware HTTP client, so it lives as long as the handler. Signed-off-by: Sivanantham Chinnaiyan --- internal/backendauth/gcp.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/internal/backendauth/gcp.go b/internal/backendauth/gcp.go index f9526c3578..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)