From f3648378ae558c77a8d1c28a87c066a1feb3b601 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Fri, 31 Jul 2026 16:18:46 +0200 Subject: [PATCH 1/6] feat(cache): optimize provider cache economics --- docs/features/cache.mdx | 23 ++ docs/pro.mdx | 20 +- internal/core/types.go | 8 + internal/providers/bedrock/chat.go | 17 +- internal/providers/cache_planner.go | 289 +++++++++++++++++++++ internal/providers/cache_planner_test.go | 71 +++++ internal/providers/gemini/gemini.go | 103 ++++++++ internal/providers/gemini/gemini_test.go | 39 +++ internal/providers/router.go | 50 +++- internal/responsecache/exact_cache_test.go | 57 ++++ internal/responsecache/simple.go | 36 ++- internal/responsecache/stream_cache.go | 22 ++ 12 files changed, 715 insertions(+), 20 deletions(-) create mode 100644 internal/providers/cache_planner.go create mode 100644 internal/providers/cache_planner_test.go diff --git a/docs/features/cache.mdx b/docs/features/cache.mdx index e44d9cb7..b150b970 100644 --- a/docs/features/cache.mdx +++ b/docs/features/cache.mdx @@ -90,6 +90,29 @@ The exact cache hashes: This means guardrails and workflows affect cache keys when they change the resolved workflow or the final body sent through execution. +JSON object member order and insignificant whitespace are canonicalized before +hashing. Requests that differ only in formatting therefore share one exact +entry. Concurrent identical misses are also coalesced in-process: one request +reaches the provider and the waiters replay its response. Failed or otherwise +non-cacheable responses are never fanned out. + +## Provider prompt-cache planning + +After routing resolves the concrete provider and model, GoModel adds a cache +plan when the stable prefix meets that provider's minimum size and the client +did not already supply a cache directive: + +- OpenAI receives a stable `prompt_cache_key`; models supporting explicit + caching also receive a breakpoint and explicit cache mode. +- Anthropic receives top-level automatic ephemeral caching. +- Amazon Bedrock Converse receives a cache point after the stable prefix. +- Native Gemini creates and reuses a short-lived cached-content object, with + concurrent creation coalesced by prefix key. + +The plan is applied post-routing so provider-only fields cannot leak through a +fallback to another backend. Cache keys include the concrete provider instance +and model, and caller-supplied cache controls always win. + ## `user_path` behavior For the exact cache, `user_path` is not added to the cache key by itself. diff --git a/docs/pro.mdx b/docs/pro.mdx index a6a9fad0..7bd85446 100644 --- a/docs/pro.mdx +++ b/docs/pro.mdx @@ -54,10 +54,16 @@ Measured on a corpus of real agent conversations at default settings: - anything, if estimated savings fall below the threshold; the request passes through byte-identical -Compression is a pure function of the request body, so the same request -always compresses the same way and the provider's prompt-cache prefix stays -stable as a conversation grows. It fails open: if the engine errors, the -request is forwarded unchanged. +Compression freezes the exact rewritten history already sent for each detected +session and replays it byte-for-byte in later requests. New turns can still be +compressed against that history, including slightly modified file reads whose +unchanged line runs become reconstructable deltas. Very long conversations roll +forward through bounded compaction epochs without revisiting older epochs. + +The final decision is economic: estimated savings on ordinary and expected +cached reuses must exceed the value of any provider-cache prefix the rewrite +would destroy. Frozen history normally makes destroyed cache value zero. The +engine fails open if it cannot safely establish that invariant. Covered endpoints: `POST /v1/chat/completions`, `POST /v1/messages`, and `POST /v1/responses` — including function and custom tool-call outputs on @@ -109,7 +115,11 @@ prompt-cache reuse. | `PRO_COMPRESSION_ENABLED` | `true` | Master switch | | `PRO_COMPRESSION_SCOPE` | `full_history` | `new_messages_only` rewrites only after the last assistant turn | | `PRO_COMPRESSION_EXCLUDE_MODELS` | — | Model globs to skip | -| `PRO_COMPRESSION_MIN_SAVINGS_TOKENS` | `64` | Minimum estimated savings to rewrite at all | +| `PRO_COMPRESSION_MIN_SAVINGS_TOKENS` | `64` | Minimum estimated savings before economic evaluation | +| `PRO_COMPRESSION_EXPECTED_CACHE_REUSES` | `4` | Future cached reuses included in expected savings | +| `PRO_COMPRESSION_CACHED_INPUT_RATE` | `0.1` | Cached-input price as a fraction of ordinary input | +| `PRO_COMPRESSION_INPUT_PRICE_PER_MTOK` | `1` | Fallback price scale for monetary audit estimates | +| `PRO_COMPRESSION_FROZEN_EPOCH_MESSAGES` | `32` | Messages per immutable compaction epoch | | `PRO_COMPRESSION_NORMALIZE` | `true` | Master switch for the normalization pass | | `PRO_COMPRESSION_STRIP_LINE_NUMBERS` | `true` | Strip the line-number gutter from file reads | | `PRO_COMPRESSION_DROP_REASONING` | `true` | Drop replayed reasoning from assistant history | diff --git a/internal/core/types.go b/internal/core/types.go index 2709e01c..96ef45a7 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -41,6 +41,14 @@ type ChatRequest struct { User string `json:"user,omitempty"` ServiceTier string `json:"service_tier,omitempty"` ExtraFields UnknownJSONFields `json:"-" swaggerignore:"true"` + // PromptCachePlan carries gateway-internal, post-routing cache metadata. + // It is never serialized to clients or upstream OpenAI-compatible APIs. + PromptCachePlan *PromptCachePlan `json:"-" swaggerignore:"true"` +} + +// PromptCachePlan identifies a provider-native cached prefix materialization. +type PromptCachePlan struct { + Key string } func (r *ChatRequest) semanticSelector() (string, string) { diff --git a/internal/providers/bedrock/chat.go b/internal/providers/bedrock/chat.go index 9fd6e3a1..d15a0703 100644 --- a/internal/providers/bedrock/chat.go +++ b/internal/providers/bedrock/chat.go @@ -153,6 +153,9 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b continue } system = append(system, &brtypes.SystemContentBlockMemberText{Value: text}) + if isGatewayCachePoint(msg.ExtraFields) { + system = append(system, &brtypes.SystemContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}) + } case "tool": block, err := convertToolResultMessage(msg) if err != nil { @@ -164,13 +167,19 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b if text == "" { continue } - appendOrMerge(brtypes.ConversationRoleUser, - []brtypes.ContentBlock{&brtypes.ContentBlockMemberText{Value: text}}) + blocks := []brtypes.ContentBlock{&brtypes.ContentBlockMemberText{Value: text}} + if isGatewayCachePoint(msg.ExtraFields) { + blocks = append(blocks, &brtypes.ContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}) + } + appendOrMerge(brtypes.ConversationRoleUser, blocks) case "assistant": blocks, err := convertAssistantMessage(msg) if err != nil { return nil, nil, err } + if isGatewayCachePoint(msg.ExtraFields) { + blocks = append(blocks, &brtypes.ContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}) + } appendOrMerge(brtypes.ConversationRoleAssistant, blocks) default: return nil, nil, core.NewInvalidRequestError("unsupported message role: "+msg.Role, nil) @@ -181,6 +190,10 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b return system, out, nil } +func isGatewayCachePoint(fields core.UnknownJSONFields) bool { + return strings.TrimSpace(string(fields.Lookup("_gomodel_cache_point"))) == "true" +} + func convertAssistantMessage(msg core.Message) ([]brtypes.ContentBlock, error) { blocks := make([]brtypes.ContentBlock, 0, 1+len(msg.ToolCalls)) if text := core.ExtractTextContent(msg.Content); text != "" { diff --git a/internal/providers/cache_planner.go b/internal/providers/cache_planner.go new file mode 100644 index 00000000..a0819dfd --- /dev/null +++ b/internal/providers/cache_planner.go @@ -0,0 +1,289 @@ +package providers + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + + "github.com/goccy/go-json" + + "github.com/enterpilot/gomodel/internal/core" +) + +const cachePointField = "_gomodel_cache_point" + +type cachePlanner struct{} + +func newCachePlanner() *cachePlanner { return &cachePlanner{} } + +func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, selector core.ModelSelector) *core.ChatRequest { + if req == nil || len(req.Messages) < 2 || hasCacheDirective(req.ExtraFields) { + return req + } + prefixBody, err := json.Marshal(struct { + Tools []map[string]any `json:"tools,omitempty"` + Messages []core.Message `json:"messages"` + }{req.Tools, req.Messages[:len(req.Messages)-1]}) + if err != nil || estimatedTokens(prefixBody) < providerCacheMinimum(providerType, selector.Model) { + return req + } + + planned, ok := cloneChatRequest(req) + if !ok { + return req + } + key := cacheAffinityKey(providerType, selector, req.User, prefixBody) + switch normalizedProviderType(providerType) { + case "openai": + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "prompt_cache_key": jsonString(key), + }) + if supportsExplicitOpenAICache(selector.Model) { + if markOpenAIChatBreakpoint(&planned.Messages) { + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "prompt_cache_options": json.RawMessage(`{"mode":"explicit"}`), + }) + } + } + case "anthropic": + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "cache_control": json.RawMessage(`{"type":"ephemeral"}`), + }) + case "bedrock": + markLastChatPrefix(&planned.Messages, cachePointField, json.RawMessage(`true`)) + case "gemini", "vertex": + planned.PromptCachePlan = &core.PromptCachePlan{Key: key} + } + return planned +} + +func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType string, selector core.ModelSelector) *core.ResponsesRequest { + if req == nil { + return req + } + items, ok := req.Input.([]core.ResponsesInputElement) + if !ok || len(items) < 2 || hasCacheDirective(req.ExtraFields) { + return req + } + prefixBody, err := json.Marshal(struct { + Instructions string `json:"instructions,omitempty"` + Tools []map[string]any `json:"tools,omitempty"` + Input []core.ResponsesInputElement `json:"input"` + }{req.Instructions, req.Tools, items[:len(items)-1]}) + if err != nil || estimatedTokens(prefixBody) < providerCacheMinimum(providerType, selector.Model) { + return req + } + planned, ok := cloneResponsesRequest(req) + if !ok { + return req + } + key := cacheAffinityKey(providerType, selector, req.User, prefixBody) + switch normalizedProviderType(providerType) { + case "openai": + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "prompt_cache_key": jsonString(key), + }) + if supportsExplicitOpenAICache(selector.Model) && markOpenAIResponsesBreakpoint(planned) { + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "prompt_cache_options": json.RawMessage(`{"mode":"explicit"}`), + }) + } + case "anthropic": + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "cache_control": json.RawMessage(`{"type":"ephemeral"}`), + }) + } + return planned +} + +func cloneChatRequest(req *core.ChatRequest) (*core.ChatRequest, bool) { + body, err := json.Marshal(req) + if err != nil { + return nil, false + } + var clone core.ChatRequest + if err := json.Unmarshal(body, &clone); err != nil { + return nil, false + } + return &clone, true +} + +func cloneResponsesRequest(req *core.ResponsesRequest) (*core.ResponsesRequest, bool) { + body, err := json.Marshal(req) + if err != nil { + return nil, false + } + var clone core.ResponsesRequest + if err := json.Unmarshal(body, &clone); err != nil { + return nil, false + } + return &clone, true +} + +func hasCacheDirective(fields core.UnknownJSONFields) bool { + for _, key := range []string{"cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options"} { + if len(fields.Lookup(key)) > 0 { + return true + } + } + return false +} + +func markLastChatPrefix(messages *[]core.Message, field string, value json.RawMessage) { + for i := len(*messages) - 2; i >= 0; i-- { + msg := &(*messages)[i] + if strings.TrimSpace(core.ExtractTextContent(msg.Content)) == "" && len(msg.ToolCalls) == 0 { + continue + } + msg.ExtraFields = mergeCacheExtras(msg.ExtraFields, map[string]json.RawMessage{field: value}) + return + } +} + +func markOpenAIChatBreakpoint(messages *[]core.Message) bool { + for i := len(*messages) - 2; i >= 0; i-- { + msg := &(*messages)[i] + switch content := msg.Content.(type) { + case string: + if content == "" { + continue + } + msg.Content = []core.ContentPart{{ + Type: "text", Text: content, + ExtraFields: core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{ + "prompt_cache_breakpoint": json.RawMessage(`{"mode":"explicit"}`), + }), + }} + return true + case []core.ContentPart: + for j := len(content) - 1; j >= 0; j-- { + if content[j].Type == "text" || content[j].Type == "image_url" || content[j].Type == "input_audio" { + content[j].ExtraFields = mergeCacheExtras(content[j].ExtraFields, map[string]json.RawMessage{ + "prompt_cache_breakpoint": json.RawMessage(`{"mode":"explicit"}`), + }) + msg.Content = content + return true + } + } + } + } + return false +} + +func markOpenAIResponsesBreakpoint(req *core.ResponsesRequest) bool { + items, ok := req.Input.([]core.ResponsesInputElement) + if !ok { + return false + } + for i := len(items) - 2; i >= 0; i-- { + switch content := items[i].Content.(type) { + case string: + if content == "" { + continue + } + items[i].Content = []core.ContentPart{{ + Type: "input_text", Text: content, + ExtraFields: core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{ + "prompt_cache_breakpoint": json.RawMessage(`{"mode":"explicit"}`), + }), + }} + req.Input = items + return true + case []core.ContentPart: + for j := len(content) - 1; j >= 0; j-- { + if content[j].Type == "input_text" || content[j].Type == "input_image" || content[j].Type == "input_file" { + content[j].ExtraFields = mergeCacheExtras(content[j].ExtraFields, map[string]json.RawMessage{ + "prompt_cache_breakpoint": json.RawMessage(`{"mode":"explicit"}`), + }) + items[i].Content = content + req.Input = items + return true + } + } + case []any: + for j := len(content) - 1; j >= 0; j-- { + block, ok := content[j].(map[string]any) + if !ok { + continue + } + blockType, _ := block["type"].(string) + if blockType == "input_text" || blockType == "input_image" || blockType == "input_file" { + if _, exists := block["prompt_cache_breakpoint"]; !exists { + block["prompt_cache_breakpoint"] = map[string]any{"mode": "explicit"} + } + items[i].Content = content + req.Input = items + return true + } + } + } + } + return false +} + +func mergeCacheExtras(base core.UnknownJSONFields, values map[string]json.RawMessage) core.UnknownJSONFields { + additions := make(map[string]json.RawMessage, len(values)) + for key, value := range values { + if len(base.Lookup(key)) == 0 { + additions[key] = value + } + } + merged, err := core.MergeUnknownJSONFields(base, additions) + if err != nil { + return base + } + return merged +} + +func cacheAffinityKey(providerType string, selector core.ModelSelector, user string, prefix []byte) string { + hash := sha256.New() + hash.Write([]byte(normalizedProviderType(providerType))) + hash.Write([]byte{0}) + hash.Write([]byte(selector.Provider)) + hash.Write([]byte{0}) + hash.Write([]byte(selector.Model)) + hash.Write([]byte{0}) + hash.Write([]byte(user)) + hash.Write([]byte{0}) + hash.Write(prefix) + return "gomodel-" + hex.EncodeToString(hash.Sum(nil)[:16]) +} + +func estimatedTokens(body []byte) int { return (len(body) + 3) / 4 } + +func providerCacheMinimum(providerType, model string) int { + providerType = normalizedProviderType(providerType) + model = strings.ToLower(model) + switch providerType { + case "openai": + return 1024 + case "anthropic": + if strings.Contains(model, "haiku-3") && !strings.Contains(model, "3-5") && !strings.Contains(model, "3.5") { + return 4096 + } + if strings.Contains(model, "haiku") { + return 2048 + } + return 1024 + case "gemini", "vertex": + return 4096 + case "bedrock": + return 1024 + default: + return int(^uint(0) >> 1) + } +} + +func normalizedProviderType(providerType string) string { + return strings.ToLower(strings.TrimSpace(providerType)) +} + +func supportsExplicitOpenAICache(model string) bool { + model = strings.ToLower(strings.TrimSpace(model)) + return strings.Contains(model, "gpt-5.6") +} + +func jsonString(value string) json.RawMessage { + body, _ := json.Marshal(value) + return body +} diff --git a/internal/providers/cache_planner_test.go b/internal/providers/cache_planner_test.go new file mode 100644 index 00000000..d7cc9a39 --- /dev/null +++ b/internal/providers/cache_planner_test.go @@ -0,0 +1,71 @@ +package providers + +import ( + "strings" + "testing" + + "github.com/goccy/go-json" + + "github.com/enterpilot/gomodel/internal/core" +) + +func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *testing.T) { + prefix := strings.Repeat("stable context ", 1500) + tests := []struct { + provider string + model string + field string + marker string + }{ + {provider: "openai", model: "gpt-5.6", field: "prompt_cache_key"}, + {provider: "anthropic", model: "claude-sonnet-4-5", field: "cache_control"}, + {provider: "bedrock", model: "anthropic.claude-sonnet-4-5", marker: cachePointField}, + {provider: "gemini", model: "gemini-2.5-pro"}, + } + for _, tt := range tests { + t.Run(tt.provider, func(t *testing.T) { + req := &core.ChatRequest{Model: tt.model, Messages: []core.Message{ + {Role: "system", Content: prefix}, + {Role: "user", Content: "new turn"}, + }} + planned := newCachePlanner().planChat(req, tt.provider, core.ModelSelector{Provider: tt.provider + "-primary", Model: tt.model}) + if planned == req { + t.Fatal("planner returned caller-owned request") + } + if tt.field != "" && len(planned.ExtraFields.Lookup(tt.field)) == 0 { + t.Fatalf("planned request lacks %q", tt.field) + } + if tt.marker != "" && len(planned.Messages[0].ExtraFields.Lookup(tt.marker)) == 0 { + t.Fatalf("stable prefix lacks %q", tt.marker) + } + if tt.provider == "gemini" && (planned.PromptCachePlan == nil || planned.PromptCachePlan.Key == "") { + t.Fatal("Gemini plan lacks an internal cached-content key") + } + if tt.provider == "openai" { + parts, ok := planned.Messages[0].Content.([]core.ContentPart) + if !ok || len(parts) != 1 || len(parts[0].ExtraFields.Lookup("prompt_cache_breakpoint")) == 0 { + t.Fatalf("OpenAI stable content lacks a breakpoint: %#v", planned.Messages[0].Content) + } + } + if !req.ExtraFields.IsEmpty() || !req.Messages[0].ExtraFields.IsEmpty() { + t.Fatal("planner mutated caller-owned request") + } + }) + } +} + +func TestCachePlannerHonorsMinimumAndClientDirective(t *testing.T) { + planner := newCachePlanner() + short := &core.ChatRequest{Messages: []core.Message{{Role: "system", Content: "short"}, {Role: "user", Content: "turn"}}} + if got := planner.planChat(short, "openai", core.ModelSelector{Model: "gpt-5.6"}); got != short { + t.Fatal("planned prefix below provider minimum") + } + + directed := &core.ChatRequest{ + Messages: []core.Message{{Role: "system", Content: strings.Repeat("x", 9000)}, {Role: "user", Content: "turn"}}, + ExtraFields: core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{"prompt_cache_key": json.RawMessage(`"client"`)}), + } + if got := planner.planChat(directed, "openai", core.ModelSelector{Model: "gpt-5.6"}); got != directed { + t.Fatal("overrode client cache directive") + } +} diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index b2a33664..9474c53c 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -10,9 +10,11 @@ import ( "os" "slices" "strings" + "sync" "time" "github.com/goccy/go-json" + "golang.org/x/sync/singleflight" "github.com/enterpilot/gomodel/internal/core" "github.com/enterpilot/gomodel/internal/httpclient" @@ -74,6 +76,14 @@ type Provider struct { useNativeAPI bool modelsURL string configErr error + cacheMu sync.Mutex + cacheObjects map[string]geminiCacheObject + cacheFlight singleflight.Group +} + +type geminiCacheObject struct { + name string + expiresAt time.Time } // New creates a new Gemini provider. @@ -462,6 +472,7 @@ func (p *Provider) nativeChatCompletion(ctx context.Context, req *core.ChatReque if err != nil { return nil, err } + p.prepareCachedContent(ctx, req, body) var geminiResp geminiGenerateContentResponse err = p.nativeClient.Do(ctx, llmclient.Request{ Method: http.MethodPost, @@ -509,6 +520,7 @@ func (p *Provider) nativeStreamChatCompletion(ctx context.Context, req *core.Cha if err != nil { return nil, err } + p.prepareCachedContent(ctx, req, body) stream, err := p.nativeClient.DoStream(ctx, llmclient.Request{ Method: http.MethodPost, Endpoint: nativeStreamEndpoint(req.Model), @@ -521,6 +533,97 @@ func (p *Provider) nativeStreamChatCompletion(ctx context.Context, req *core.Cha return newGeminiNativeStream(stream, req.Model, includeUsage, p.responseProviderName()), nil } +type geminiCreateCachedContentRequest struct { + Model string `json:"model"` + SystemInstruction *geminiContent `json:"systemInstruction,omitempty"` + Contents []geminiContent `json:"contents"` + Tools []geminiTool `json:"tools,omitempty"` + TTL string `json:"ttl"` +} + +type geminiCreateCachedContentResponse struct { + Name string `json:"name"` + ExpireTime time.Time `json:"expireTime"` +} + +// prepareCachedContent materializes the stable prefix selected by the +// post-routing planner. Cache creation is best-effort: an unsupported model or +// endpoint must never turn a request that would otherwise work into a failure. +func (p *Provider) prepareCachedContent(ctx context.Context, req *core.ChatRequest, body *geminiGenerateContentRequest) { + if body == nil || body.CachedContent != "" || len(body.Contents) < 2 { + return + } + if req.PromptCachePlan == nil || strings.TrimSpace(req.PromptCachePlan.Key) == "" { + return + } + key := req.PromptCachePlan.Key + if cached := p.cachedContentObject(key); cached != "" { + useGeminiCachedPrefix(body, cached) + return + } + + value, err, _ := p.cacheFlight.Do(key, func() (any, error) { + if cached := p.cachedContentObject(key); cached != "" { + return cached, nil + } + createReq := geminiCreateCachedContentRequest{ + Model: "models/" + normalizeGeminiModelID(req.Model), + SystemInstruction: body.SystemInstruction, + Contents: append([]geminiContent(nil), body.Contents[:len(body.Contents)-1]...), + Tools: append([]geminiTool(nil), body.Tools...), + TTL: "300s", + } + var created geminiCreateCachedContentResponse + if err := p.nativeClient.Do(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: "/cachedContents", + Body: &createReq, + }, &created); err != nil { + return "", err + } + if strings.TrimSpace(created.Name) == "" { + return "", fmt.Errorf("cached-content creation for Gemini returned an empty name") + } + expiresAt := created.ExpireTime + if expiresAt.IsZero() { + expiresAt = time.Now().Add(5 * time.Minute) + } + p.cacheMu.Lock() + if p.cacheObjects == nil { + p.cacheObjects = make(map[string]geminiCacheObject) + } + p.cacheObjects[key] = geminiCacheObject{name: created.Name, expiresAt: expiresAt} + p.cacheMu.Unlock() + return created.Name, nil + }) + if err == nil { + if cached, ok := value.(string); ok && cached != "" { + useGeminiCachedPrefix(body, cached) + } + } +} + +func (p *Provider) cachedContentObject(key string) string { + p.cacheMu.Lock() + defer p.cacheMu.Unlock() + entry, ok := p.cacheObjects[key] + if !ok { + return "" + } + if time.Now().Add(15 * time.Second).After(entry.expiresAt) { + delete(p.cacheObjects, key) + return "" + } + return entry.name +} + +func useGeminiCachedPrefix(body *geminiGenerateContentRequest, name string) { + body.CachedContent = name + body.Contents = append([]geminiContent(nil), body.Contents[len(body.Contents)-1]) + body.SystemInstruction = nil + body.Tools = nil +} + // geminiModel represents a model in Gemini's native API response type geminiModel struct { Name string `json:"name"` diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index f9de0d4f..9826b721 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync/atomic" "testing" "github.com/enterpilot/gomodel/internal/core" @@ -34,6 +35,44 @@ func TestNew(t *testing.T) { } } +func TestPrepareCachedContentCreatesAndReusesObject(t *testing.T) { + var creates atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/cachedContents" { + t.Fatalf("path = %q, want /cachedContents", r.URL.Path) + } + creates.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"name":"cachedContents/session-prefix","expireTime":"2099-01-01T00:00:00Z"}`) + })) + defer server.Close() + + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) + req := &core.ChatRequest{Model: "gemini-2.5-pro", PromptCachePlan: &core.PromptCachePlan{Key: "prefix-key"}} + newBody := func() *geminiGenerateContentRequest { + return &geminiGenerateContentRequest{ + SystemInstruction: &geminiContent{Parts: []geminiPart{{Text: "system"}}}, + Contents: []geminiContent{ + {Role: "user", Parts: []geminiPart{{Text: "stable"}}}, + {Role: "user", Parts: []geminiPart{{Text: "dynamic"}}}, + }, + } + } + first := newBody() + p.prepareCachedContent(context.Background(), req, first) + second := newBody() + p.prepareCachedContent(context.Background(), req, second) + if creates.Load() != 1 { + t.Fatalf("cache creates = %d, want 1", creates.Load()) + } + for i, body := range []*geminiGenerateContentRequest{first, second} { + if body.CachedContent != "cachedContents/session-prefix" || len(body.Contents) != 1 || body.SystemInstruction != nil { + t.Fatalf("body %d did not use cached prefix: %+v", i, body) + } + } +} + func TestNew_ReturnsProvider(t *testing.T) { provider := New(providers.ProviderConfig{APIKey: "test-api-key"}, providers.ProviderOptions{}) diff --git a/internal/providers/router.go b/internal/providers/router.go index bc20616d..e5b814c0 100644 --- a/internal/providers/router.go +++ b/internal/providers/router.go @@ -21,7 +21,8 @@ var ErrRegistryNotInitialized = fmt.Errorf("model registry has no models: ensure // It uses a dynamic model-to-provider mapping that is populated at startup // by fetching available models from each provider's /models endpoint. type Router struct { - lookup core.ModelLookup + lookup core.ModelLookup + cachePlanner *cachePlanner } type providerTypeRegistry interface { @@ -76,7 +77,8 @@ func NewRouter(lookup core.ModelLookup) (*Router, error) { return nil, fmt.Errorf("lookup cannot be nil") } return &Router{ - lookup: lookup, + lookup: lookup, + cachePlanner: newCachePlanner(), }, nil } @@ -478,6 +480,17 @@ func routeStampedModelResponse[Req any, Resp any]( return stampProvider(resp, providerType), nil } +func routeModelStream[Req any]( + r *Router, + ctx context.Context, + model, providerHint string, + buildForward func(core.ModelSelector) Req, + call func(context.Context, core.Provider, Req) (io.ReadCloser, error), +) (io.ReadCloser, error) { + stream, _, err := routeResolvedModelCall(r, ctx, model, providerHint, buildForward, call) + return stream, err +} + func routeNativeBatchCall[T any](r *Router, ctx context.Context, providerType string, call func(context.Context, core.NativeBatchProvider) (T, error)) (T, error) { bp, err := r.resolveNativeBatchProvider(providerType) if err != nil { @@ -567,6 +580,25 @@ func forwardResponsesRequest(req *core.ResponsesRequest, selector core.ModelSele return &forwardReq } +func (r *Router) plannedChatRequest(ctx context.Context, req *core.ChatRequest, selector core.ModelSelector) *core.ChatRequest { + forward := r.forwardChatRequest(ctx, req, selector) + if r.cachePlanner == nil || len(forward.Messages) < 2 { + return forward + } + return r.cachePlanner.planChat(forward, r.lookup.GetProviderType(selector.QualifiedModel()), selector) +} + +func (r *Router) plannedResponsesRequest(req *core.ResponsesRequest, selector core.ModelSelector) *core.ResponsesRequest { + forward := forwardResponsesRequest(req, selector) + if r.cachePlanner == nil { + return forward + } + if items, ok := forward.Input.([]core.ResponsesInputElement); !ok || len(items) < 2 { + return forward + } + return r.cachePlanner.planResponses(forward, r.lookup.GetProviderType(selector.QualifiedModel()), selector) +} + func forwardEmbeddingRequest(req *core.EmbeddingRequest, selector core.ModelSelector) *core.EmbeddingRequest { forwardReq := *req forwardReq.Model = selector.Model @@ -627,7 +659,7 @@ func (r *Router) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*co req.Model, req.Provider, func(selector core.ModelSelector) *core.ChatRequest { - return r.forwardChatRequest(ctx, req, selector) + return r.plannedChatRequest(ctx, req, selector) }, callChatCompletion, ) @@ -636,19 +668,18 @@ func (r *Router) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*co // StreamChatCompletion routes the streaming request to the appropriate provider. // Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { - stream, _, err := routeResolvedModelCall( + return routeModelStream( r, ctx, req.Model, req.Provider, func(selector core.ModelSelector) *core.ChatRequest { - return r.forwardChatRequest(ctx, req, selector) + return r.plannedChatRequest(ctx, req, selector) }, func(ctx context.Context, provider core.Provider, forwardReq *core.ChatRequest) (io.ReadCloser, error) { return provider.StreamChatCompletion(ctx, forwardReq) }, ) - return stream, err } // ListModels returns all models from the lookup. @@ -678,7 +709,7 @@ func (r *Router) Responses(ctx context.Context, req *core.ResponsesRequest) (*co req.Model, req.Provider, func(selector core.ModelSelector) *core.ResponsesRequest { - return forwardResponsesRequest(req, selector) + return r.plannedResponsesRequest(req, selector) }, callResponses, ) @@ -687,19 +718,18 @@ func (r *Router) Responses(ctx context.Context, req *core.ResponsesRequest) (*co // StreamResponses routes the streaming Responses API request to the appropriate provider. // Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { - stream, _, err := routeResolvedModelCall( + return routeModelStream( r, ctx, req.Model, req.Provider, func(selector core.ModelSelector) *core.ResponsesRequest { - return forwardResponsesRequest(req, selector) + return r.plannedResponsesRequest(req, selector) }, func(ctx context.Context, provider core.Provider, forwardReq *core.ResponsesRequest) (io.ReadCloser, error) { return provider.StreamResponses(ctx, forwardReq) }, ) - return stream, err } // Embeddings routes the embeddings request to the appropriate provider. diff --git a/internal/responsecache/exact_cache_test.go b/internal/responsecache/exact_cache_test.go index a934bd93..c4956559 100644 --- a/internal/responsecache/exact_cache_test.go +++ b/internal/responsecache/exact_cache_test.go @@ -165,6 +165,63 @@ func TestHandleRequest_DifferentBodyDifferentKey(t *testing.T) { } } +func TestHandleRequest_CoalescesConcurrentIdenticalMisses(t *testing.T) { + store := cache.NewMapStore() + defer store.Close() + mw := NewResponseCacheMiddlewareWithStore(store, time.Hour) + workflow := resolvedWorkflow("openai", "gpt-4") + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) + + var calls atomic.Int32 + started := make(chan struct{}) + release := make(chan struct{}) + next := func(c *echo.Context) error { + if calls.Add(1) == 1 { + close(started) + } + <-release + return c.JSON(http.StatusOK, map[string]string{"result": "shared"}) + } + + const requests = 12 + recorders := make([]*httptest.ResponseRecorder, requests) + var wg sync.WaitGroup + for i := range requests { + wg.Go(func() { + recorders[i] = driveHandleRequest(t, mw, workflow, body, nil, next) + }) + } + <-started + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + if got := calls.Load(); got != 1 { + t.Fatalf("provider calls = %d, want one coalesced miss", got) + } + hits := 0 + for i, rec := range recorders { + if rec.Code != http.StatusOK || !bytes.Contains(rec.Body.Bytes(), []byte("shared")) { + t.Fatalf("response %d = status %d body %q", i, rec.Code, rec.Body.String()) + } + if rec.Header().Get("X-Cache") == "HIT (exact)" { + hits++ + } + } + if hits != requests-1 { + t.Fatalf("coalesced hit responses = %d, want %d", hits, requests-1) + } +} + +func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { + plan := resolvedWorkflow("openai", "gpt-4") + compact := []byte(`{"input":[1,2],"model":"gpt-4"}`) + formatted := []byte("{\n \"model\": \"gpt-4\",\n \"input\": [1, 2]\n}") + if first, second := hashRequest("/v1/embeddings", compact, plan), hashRequest("/v1/embeddings", formatted, plan); first != second { + t.Fatalf("canonical-equivalent JSON produced different keys: %s != %s", first, second) + } +} + func TestHashRequest_ResolvedModelChangesKey(t *testing.T) { body := []byte(`{"model":"anthropic/claude-opus-4-6","messages":[{"role":"user","content":"hi"}]}`) diff --git a/internal/responsecache/simple.go b/internal/responsecache/simple.go index e2cd176f..c23aad13 100644 --- a/internal/responsecache/simple.go +++ b/internal/responsecache/simple.go @@ -15,6 +15,7 @@ import ( "github.com/labstack/echo/v5" "github.com/tidwall/gjson" + "golang.org/x/sync/singleflight" "github.com/enterpilot/gomodel/internal/cache" "github.com/enterpilot/gomodel/internal/core" @@ -47,6 +48,7 @@ type simpleCacheMiddleware struct { workers sync.WaitGroup mu sync.RWMutex closed bool + misses singleflight.Group } func newSimpleCacheMiddleware(store cache.Store, ttl time.Duration, hitRecorder func(exchange, []byte, string)) *simpleCacheMiddleware { @@ -101,14 +103,42 @@ func (m *simpleCacheMiddleware) StoreAfter(ex exchange, body []byte, next func() plan := core.GetWorkflow(ex.Context()) key := hashRequest(path, body, plan) - data, ok, err := ex.Capture("response cache: failed to capture cacheable response body", next) + type missResult struct { + owner *struct{ marker byte } + data []byte + } + owner := &struct{ marker byte }{} + value, err, _ := m.misses.Do(key, func() (any, error) { + data, ok, err := ex.Capture("response cache: failed to capture cacheable response body", next) + if err != nil { + return nil, err + } + if !ok { + return missResult{owner: owner}, nil + } + m.enqueueWrite(cacheWriteJob{key: key, data: data}) + return missResult{owner: owner, data: data}, nil + }) if err != nil { return err } - if !ok { + result, _ := value.(missResult) + if result.owner == owner { return nil } - m.enqueueWrite(cacheWriteJob{key: key, data: data}) + // The leader produced a non-cacheable result (failure status, failover, or + // malformed body). Waiting followers must execute independently rather than + // replaying something the normal cache would refuse to store. + if len(result.data) == 0 { + return next() + } + if err := ex.ReplayHit(body, result.data, CacheTypeExact); err != nil { + return next() + } + ex.MarkHit(CacheTypeExact) + if m.hitRecorder != nil { + m.hitRecorder(ex, result.data, CacheTypeExact) + } return nil } diff --git a/internal/responsecache/stream_cache.go b/internal/responsecache/stream_cache.go index 6353fd1e..68ba3454 100644 --- a/internal/responsecache/stream_cache.go +++ b/internal/responsecache/stream_cache.go @@ -2,6 +2,7 @@ package responsecache import ( "bytes" + "io" "net/http" "strings" @@ -53,8 +54,29 @@ func cacheKeyRequestBody(path string, body []byte) []byte { } return normalized default: + return canonicalJSONForCache(body) + } +} + +// canonicalJSONForCache makes semantically identical JSON bodies share an +// exact-cache key despite insignificant whitespace or object-key ordering. It +// preserves number spellings through json.Number and falls back byte-for-byte +// for malformed or multi-value input. +func canonicalJSONForCache(body []byte) []byte { + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return body + } + if err := decoder.Decode(new(any)); err != io.EOF { + return body + } + canonical, err := json.Marshal(value) + if err != nil { return body } + return canonical } func isEventStreamContentType(contentType string) bool { From 674403170b1b68b073d064f92ce90849c9360a20 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sat, 1 Aug 2026 13:21:49 +0200 Subject: [PATCH 2/6] feat(cache): add provider planner kill switch --- .env.template | 5 +++++ docs/features/cache.mdx | 8 +++++++ internal/providers/cache_planner.go | 27 ++++++++++++++++++++++-- internal/providers/cache_planner_test.go | 20 ++++++++++++++++++ 4 files changed, 58 insertions(+), 2 deletions(-) diff --git a/.env.template b/.env.template index 1605c1fc..b1cecc0a 100644 --- a/.env.template +++ b/.env.template @@ -109,6 +109,11 @@ # Model cache uses the local filesystem by default. # Set REDIS_URL to use Redis-backed caching instead. +# Add provider-native prompt-cache keys, breakpoints, and cached-content plans +# after routing (default: true). Set false as an operational kill switch; client +# cache directives still pass through according to the selected provider. +# PROVIDER_PROMPT_CACHE_PLANNER_ENABLED=true + # Redis Configuration # REDIS_URL=redis://localhost:6379 # REDIS_KEY_MODELS=gomodel:models diff --git a/docs/features/cache.mdx b/docs/features/cache.mdx index b150b970..3f1996f8 100644 --- a/docs/features/cache.mdx +++ b/docs/features/cache.mdx @@ -113,6 +113,14 @@ The plan is applied post-routing so provider-only fields cannot leak through a fallback to another backend. Cache keys include the concrete provider instance and model, and caller-supplied cache controls always win. +Provider prompt-cache planning is enabled by default. To disable GoModel's +automatic planning while continuing to pass through compatible client cache +directives, set: + +```bash +PROVIDER_PROMPT_CACHE_PLANNER_ENABLED=false +``` + ## `user_path` behavior For the exact cache, `user_path` is not added to the cache key by itself. diff --git a/internal/providers/cache_planner.go b/internal/providers/cache_planner.go index a0819dfd..7f2a9072 100644 --- a/internal/providers/cache_planner.go +++ b/internal/providers/cache_planner.go @@ -3,6 +3,9 @@ package providers import ( "crypto/sha256" "encoding/hex" + "log/slog" + "os" + "strconv" "strings" "github.com/goccy/go-json" @@ -10,11 +13,31 @@ import ( "github.com/enterpilot/gomodel/internal/core" ) -const cachePointField = "_gomodel_cache_point" +const ( + cachePointField = "_gomodel_cache_point" + providerPromptCachePlannerEnabledEnv = "PROVIDER_PROMPT_CACHE_PLANNER_ENABLED" +) type cachePlanner struct{} -func newCachePlanner() *cachePlanner { return &cachePlanner{} } +func newCachePlanner() *cachePlanner { + raw, configured := os.LookupEnv(providerPromptCachePlannerEnabledEnv) + if !configured || strings.TrimSpace(raw) == "" { + return &cachePlanner{} + } + enabled, err := strconv.ParseBool(strings.TrimSpace(raw)) + if err != nil { + slog.Warn("invalid provider prompt-cache planner flag; using default", + "env", providerPromptCachePlannerEnabledEnv, + "value", raw, + "default", true) + return &cachePlanner{} + } + if !enabled { + return nil + } + return &cachePlanner{} +} func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, selector core.ModelSelector) *core.ChatRequest { if req == nil || len(req.Messages) < 2 || hasCacheDirective(req.ExtraFields) { diff --git a/internal/providers/cache_planner_test.go b/internal/providers/cache_planner_test.go index d7cc9a39..bb0d8050 100644 --- a/internal/providers/cache_planner_test.go +++ b/internal/providers/cache_planner_test.go @@ -54,6 +54,26 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes } } +func TestNewCachePlanner_EnvironmentKillSwitch(t *testing.T) { + for _, tt := range []struct { + name string + value string + enabled bool + }{ + {name: "default on", enabled: true}, + {name: "explicit on", value: "true", enabled: true}, + {name: "explicit off", value: "false", enabled: false}, + {name: "invalid keeps safe default", value: "sometimes", enabled: true}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Setenv(providerPromptCachePlannerEnabledEnv, tt.value) + if got := newCachePlanner() != nil; got != tt.enabled { + t.Fatalf("newCachePlanner() enabled = %v, want %v", got, tt.enabled) + } + }) + } +} + func TestCachePlannerHonorsMinimumAndClientDirective(t *testing.T) { planner := newCachePlanner() short := &core.ChatRequest{Messages: []core.Message{{Role: "system", Content: "short"}, {Role: "user", Content: "turn"}}} From 0a2a57be83420b297025a1f26d4d07d93d327a3a Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sat, 1 Aug 2026 19:35:19 +0200 Subject: [PATCH 3/6] fix(cache): harden provider-aware cache planning --- .env.template | 4 +- docs/features/cache.mdx | 7 +- docs/providers/gemini.mdx | 6 + internal/core/types.go | 4 + internal/providers/bedrock/bedrock_test.go | 54 +++++ internal/providers/bedrock/chat.go | 88 +++++++- internal/providers/bedrock/chat_stream.go | 19 +- internal/providers/cache_control.go | 8 +- internal/providers/cache_planner.go | 241 ++++++++++++++++++--- internal/providers/cache_planner_test.go | 126 ++++++++++- internal/providers/gemini/gemini.go | 123 ++++++++--- internal/providers/gemini/gemini_test.go | 186 +++++++++++++++- internal/providers/keyring.go | 22 ++ internal/providers/keyring_test.go | 29 +++ internal/providers/router.go | 5 +- internal/responsecache/exact_cache_test.go | 201 +++++++++++++---- internal/responsecache/simple.go | 73 +++++-- 17 files changed, 1030 insertions(+), 166 deletions(-) diff --git a/.env.template b/.env.template index b1cecc0a..d1e98442 100644 --- a/.env.template +++ b/.env.template @@ -110,8 +110,8 @@ # Set REDIS_URL to use Redis-backed caching instead. # Add provider-native prompt-cache keys, breakpoints, and cached-content plans -# after routing (default: true). Set false as an operational kill switch; client -# cache directives still pass through according to the selected provider. +# after routing (default: true). Set false as an operational kill switch; empty +# or invalid values keep the default. Client cache directives still pass through. # PROVIDER_PROMPT_CACHE_PLANNER_ENABLED=true # Redis Configuration diff --git a/docs/features/cache.mdx b/docs/features/cache.mdx index 3f1996f8..2b66e14a 100644 --- a/docs/features/cache.mdx +++ b/docs/features/cache.mdx @@ -106,8 +106,8 @@ did not already supply a cache directive: caching also receive a breakpoint and explicit cache mode. - Anthropic receives top-level automatic ephemeral caching. - Amazon Bedrock Converse receives a cache point after the stable prefix. -- Native Gemini creates and reuses a short-lived cached-content object, with - concurrent creation coalesced by prefix key. +- Native Gemini AI Studio creates and reuses a five-minute cached-content + object, with concurrent creation coalesced by prefix and stable credential. The plan is applied post-routing so provider-only fields cannot leak through a fallback to another backend. Cache keys include the concrete provider instance @@ -121,6 +121,9 @@ directives, set: PROVIDER_PROMPT_CACHE_PLANNER_ENABLED=false ``` +The flag accepts Go boolean values. Missing, empty, or invalid values keep the +safe default (`true`) and an invalid value emits a startup warning. + ## `user_path` behavior For the exact cache, `user_path` is not added to the cache key by itself. diff --git a/docs/providers/gemini.mdx b/docs/providers/gemini.mdx index 658703e8..b3f1c6cd 100644 --- a/docs/providers/gemini.mdx +++ b/docs/providers/gemini.mdx @@ -35,6 +35,12 @@ providers: when per-provider `GEMINI_API_MODE` is unset. Prefer `GEMINI_API_MODE`. +Native AI Studio mode can also materialize GoModel's post-routing prompt-cache +plans as five-minute Gemini cached-content objects. Reuse is scoped to the API +key selected for the session; sessionless traffic with rotating keys skips +object creation so a resource is never reused under the wrong credential. +Vertex currently receives the original request without cached-content objects. + ## Base URLs GoModel keeps separate internal bases for native Gemini and the diff --git a/internal/core/types.go b/internal/core/types.go index 96ef45a7..b4ec8e45 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -46,6 +46,10 @@ type ChatRequest struct { PromptCachePlan *PromptCachePlan `json:"-" swaggerignore:"true"` } +// GatewayCachePointField is the internal marker shared by provider cache +// planners and native request translators. It is never forwarded verbatim. +const GatewayCachePointField = "_gomodel_cache_point" + // PromptCachePlan identifies a provider-native cached prefix materialization. type PromptCachePlan struct { Key string diff --git a/internal/providers/bedrock/bedrock_test.go b/internal/providers/bedrock/bedrock_test.go index c946d691..8d4595c7 100644 --- a/internal/providers/bedrock/bedrock_test.go +++ b/internal/providers/bedrock/bedrock_test.go @@ -798,3 +798,57 @@ func TestStreamConverter_FormatChunkUsage(t *testing.T) { t.Errorf("total_tokens = %v", usage["total_tokens"]) } } + +func TestGatewayCachePointUsesJSONBooleanAndNeverCreatesEmptyMessage(t *testing.T) { + fields := core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{ + core.GatewayCachePointField: json.RawMessage(" true\n"), + }) + if !isGatewayCachePoint(fields) { + t.Fatal("formatted JSON true was not recognized") + } + falseFields := core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{ + core.GatewayCachePointField: json.RawMessage(`false`), + }) + if isGatewayCachePoint(falseFields) { + t.Fatal("JSON false was recognized as a cache point") + } + _, messages, err := convertMessages([]core.Message{{Role: "assistant", ExtraFields: fields}}) + if err != nil { + t.Fatal(err) + } + if len(messages) != 0 { + t.Fatalf("empty assistant created a cache-only message: %+v", messages) + } +} + +type testBedrockAPIError struct{ code, message string } + +func (e testBedrockAPIError) Error() string { return e.code + ": " + e.message } +func (e testBedrockAPIError) ErrorCode() string { return e.code } +func (e testBedrockAPIError) ErrorMessage() string { return e.message } + +func TestCachePointFallbackIsNarrowAndLossless(t *testing.T) { + if !isCachePointValidationError(testBedrockAPIError{"ValidationException", "cache point below minimum tokens"}) { + t.Fatal("cache-point validation error was not recognized") + } + if isCachePointValidationError(testBedrockAPIError{"ValidationException", "invalid tool schema"}) { + t.Fatal("unrelated validation error would trigger a retry") + } + parts := converseParts{ + system: []brtypes.SystemContentBlock{ + &brtypes.SystemContentBlockMemberText{Value: "system"}, + &brtypes.SystemContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}, + }, + messages: []brtypes.Message{{Role: brtypes.ConversationRoleUser, Content: []brtypes.ContentBlock{ + &brtypes.ContentBlockMemberText{Value: "user"}, + &brtypes.ContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}, + }}}, + } + clean := withoutCachePoints(parts) + if partsHaveCachePoints(clean) || len(clean.system) != 1 || len(clean.messages) != 1 || len(clean.messages[0].Content) != 1 { + t.Fatalf("cache-point removal damaged request content: %+v", clean) + } + if !partsHaveCachePoints(parts) { + t.Fatal("cache-point removal mutated original parts") + } +} diff --git a/internal/providers/bedrock/chat.go b/internal/providers/bedrock/chat.go index d15a0703..60bd1a01 100644 --- a/internal/providers/bedrock/chat.go +++ b/internal/providers/bedrock/chat.go @@ -2,6 +2,7 @@ package bedrock import ( "context" + "errors" "fmt" "math" "strings" @@ -32,13 +33,11 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* return nil, err } - out, err := p.runtime.Converse(ctx, &bedrockruntime.ConverseInput{ - ModelId: parts.modelID, - Messages: parts.messages, - System: parts.system, - InferenceConfig: parts.infCfg, - ToolConfig: parts.toolCfg, - }) + out, err := p.runtime.Converse(ctx, converseInput(parts)) + if err != nil && partsHaveCachePoints(parts) && isCachePointValidationError(err) { + parts = withoutCachePoints(parts) + out, err = p.runtime.Converse(ctx, converseInput(parts)) + } if err != nil { return nil, mapAWSError(err) } @@ -46,6 +45,74 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* return convertConverseOutput(req.Model, out), nil } +func converseInput(parts converseParts) *bedrockruntime.ConverseInput { + return &bedrockruntime.ConverseInput{ + ModelId: parts.modelID, Messages: parts.messages, System: parts.system, + InferenceConfig: parts.infCfg, ToolConfig: parts.toolCfg, + } +} + +func partsHaveCachePoints(parts converseParts) bool { + for _, block := range parts.system { + if _, ok := block.(*brtypes.SystemContentBlockMemberCachePoint); ok { + return true + } + } + for _, message := range parts.messages { + for _, block := range message.Content { + if _, ok := block.(*brtypes.ContentBlockMemberCachePoint); ok { + return true + } + } + } + return false +} + +// withoutCachePoints makes the planner best-effort. If a model or regional +// endpoint rejects Bedrock cache points, the request is retried once without +// gateway-generated markers rather than failing an otherwise valid call. +func withoutCachePoints(parts converseParts) converseParts { + clean := parts + clean.system = make([]brtypes.SystemContentBlock, 0, len(parts.system)) + for _, block := range parts.system { + if _, ok := block.(*brtypes.SystemContentBlockMemberCachePoint); !ok { + clean.system = append(clean.system, block) + } + } + clean.messages = make([]brtypes.Message, 0, len(parts.messages)) + for _, message := range parts.messages { + copyMessage := message + copyMessage.Content = make([]brtypes.ContentBlock, 0, len(message.Content)) + for _, block := range message.Content { + if _, ok := block.(*brtypes.ContentBlockMemberCachePoint); !ok { + copyMessage.Content = append(copyMessage.Content, block) + } + } + if len(copyMessage.Content) > 0 { + clean.messages = append(clean.messages, copyMessage) + } + } + return clean +} + +func isCachePointValidationError(err error) bool { + type apiError interface { + error + ErrorCode() string + ErrorMessage() string + } + var apiErr apiError + code, message := "", err.Error() + if errors.As(err, &apiErr) { + code, message = apiErr.ErrorCode(), apiErr.ErrorMessage() + } + code = strings.ToLower(code) + message = strings.ToLower(message) + validation := strings.Contains(code, "validation") || strings.Contains(code, "badrequest") + return validation && strings.Contains(message, "cache") && + (strings.Contains(message, "point") || strings.Contains(message, "minimum") || strings.Contains(message, "token")) +} + // converseParts holds the shared pieces of a Converse / ConverseStream request // so the two API entry points can stay otherwise independent. type converseParts struct { @@ -177,7 +244,9 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b if err != nil { return nil, nil, err } - if isGatewayCachePoint(msg.ExtraFields) { + // A cache point is a boundary after real content, never content by + // itself. Bedrock rejects cache-point-only messages. + if len(blocks) > 0 && isGatewayCachePoint(msg.ExtraFields) { blocks = append(blocks, &brtypes.ContentBlockMemberCachePoint{Value: brtypes.CachePointBlock{Type: brtypes.CachePointTypeDefault}}) } appendOrMerge(brtypes.ConversationRoleAssistant, blocks) @@ -191,7 +260,8 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b } func isGatewayCachePoint(fields core.UnknownJSONFields) bool { - return strings.TrimSpace(string(fields.Lookup("_gomodel_cache_point"))) == "true" + var enabled bool + return json.Unmarshal(fields.Lookup(core.GatewayCachePointField), &enabled) == nil && enabled } func convertAssistantMessage(msg core.Message) ([]brtypes.ContentBlock, error) { diff --git a/internal/providers/bedrock/chat_stream.go b/internal/providers/bedrock/chat_stream.go index ecbc232c..08523043 100644 --- a/internal/providers/bedrock/chat_stream.go +++ b/internal/providers/bedrock/chat_stream.go @@ -29,13 +29,11 @@ func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatReque return nil, err } - out, err := p.runtime.ConverseStream(ctx, &bedrockruntime.ConverseStreamInput{ - ModelId: parts.modelID, - Messages: parts.messages, - System: parts.system, - InferenceConfig: parts.infCfg, - ToolConfig: parts.toolCfg, - }) + out, err := p.runtime.ConverseStream(ctx, converseStreamInput(parts)) + if err != nil && partsHaveCachePoints(parts) && isCachePointValidationError(err) { + parts = withoutCachePoints(parts) + out, err = p.runtime.ConverseStream(ctx, converseStreamInput(parts)) + } if err != nil { return nil, mapAWSError(err) } @@ -43,6 +41,13 @@ func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatReque return newOpenAIStream(out, req.Model), nil } +func converseStreamInput(parts converseParts) *bedrockruntime.ConverseStreamInput { + return &bedrockruntime.ConverseStreamInput{ + ModelId: parts.modelID, Messages: parts.messages, System: parts.system, + InferenceConfig: parts.infCfg, ToolConfig: parts.toolCfg, + } +} + // streamConverter consumes a Bedrock ConverseStream event channel and emits // OpenAI-compatible SSE chunks via io.Reader. Reads are buffered: each loop // iteration consumes events until at least one chunk's worth of data is ready diff --git a/internal/providers/cache_control.go b/internal/providers/cache_control.go index a7327f05..090f7ed9 100644 --- a/internal/providers/cache_control.go +++ b/internal/providers/cache_control.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "slices" - "strings" "github.com/goccy/go-json" @@ -74,12 +73,7 @@ func adaptAnthropicBatchCacheControl(ctx context.Context, req *core.BatchRequest } func providerAcceptsAnthropicCacheControl(providerType string) bool { - switch strings.ToLower(strings.TrimSpace(providerType)) { - case "anthropic", "openrouter": - return true - default: - return false - } + return promptCacheProfileFor(providerType).acceptsAnthropicCacheControl } func hasAnthropicCacheControl(req *core.ChatRequest) bool { diff --git a/internal/providers/cache_planner.go b/internal/providers/cache_planner.go index 7f2a9072..dafa1723 100644 --- a/internal/providers/cache_planner.go +++ b/internal/providers/cache_planner.go @@ -5,6 +5,7 @@ import ( "encoding/hex" "log/slog" "os" + "slices" "strconv" "strings" @@ -14,16 +15,32 @@ import ( ) const ( - cachePointField = "_gomodel_cache_point" providerPromptCachePlannerEnabledEnv = "PROVIDER_PROMPT_CACHE_PLANNER_ENABLED" ) -type cachePlanner struct{} +type promptCacheMode uint8 + +const ( + promptCacheUnsupported promptCacheMode = iota + promptCacheOpenAI + promptCacheAnthropic + promptCacheBedrock + promptCacheGemini +) + +type promptCacheProfile struct { + mode promptCacheMode + acceptsAnthropicCacheControl bool +} + +type cachePlanner struct { + enabled bool +} func newCachePlanner() *cachePlanner { raw, configured := os.LookupEnv(providerPromptCachePlannerEnabledEnv) if !configured || strings.TrimSpace(raw) == "" { - return &cachePlanner{} + return &cachePlanner{enabled: true} } enabled, err := strconv.ParseBool(strings.TrimSpace(raw)) if err != nil { @@ -31,23 +48,29 @@ func newCachePlanner() *cachePlanner { "env", providerPromptCachePlannerEnabledEnv, "value", raw, "default", true) - return &cachePlanner{} - } - if !enabled { - return nil + return &cachePlanner{enabled: true} } - return &cachePlanner{} + return &cachePlanner{enabled: enabled} } func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, selector core.ModelSelector) *core.ChatRequest { - if req == nil || len(req.Messages) < 2 || hasCacheDirective(req.ExtraFields) { + profile := promptCacheProfileFor(providerType) + if p == nil || !p.enabled || req == nil || len(req.Messages) < 2 || profile.mode == promptCacheUnsupported { + return req + } + minimum := providerCacheMinimum(profile, selector.Model) + if tokens, conclusive := estimateSimpleChatPrefixTokens(req); conclusive && tokens < minimum { + return req + } + prefixMessages := req.Messages[:len(req.Messages)-1] + if hasChatCacheDirective(req, prefixMessages) { return req } prefixBody, err := json.Marshal(struct { Tools []map[string]any `json:"tools,omitempty"` Messages []core.Message `json:"messages"` - }{req.Tools, req.Messages[:len(req.Messages)-1]}) - if err != nil || estimatedTokens(prefixBody) < providerCacheMinimum(providerType, selector.Model) { + }{req.Tools, prefixMessages}) + if err != nil || estimatedTokens(prefixBody) < minimum { return req } @@ -56,8 +79,8 @@ func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, sele return req } key := cacheAffinityKey(providerType, selector, req.User, prefixBody) - switch normalizedProviderType(providerType) { - case "openai": + switch profile.mode { + case promptCacheOpenAI: planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ "prompt_cache_key": jsonString(key), }) @@ -68,24 +91,25 @@ func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, sele }) } } - case "anthropic": + case promptCacheAnthropic: planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ "cache_control": json.RawMessage(`{"type":"ephemeral"}`), }) - case "bedrock": - markLastChatPrefix(&planned.Messages, cachePointField, json.RawMessage(`true`)) - case "gemini", "vertex": + case promptCacheBedrock: + markLastChatPrefix(&planned.Messages, core.GatewayCachePointField, json.RawMessage(`true`)) + case promptCacheGemini: planned.PromptCachePlan = &core.PromptCachePlan{Key: key} } return planned } func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType string, selector core.ModelSelector) *core.ResponsesRequest { - if req == nil { + profile := promptCacheProfileFor(providerType) + if p == nil || !p.enabled || req == nil || profile.mode == promptCacheUnsupported { return req } items, ok := req.Input.([]core.ResponsesInputElement) - if !ok || len(items) < 2 || hasCacheDirective(req.ExtraFields) { + if !ok || len(items) < 2 || hasResponsesCacheDirective(req, items[:len(items)-1]) { return req } prefixBody, err := json.Marshal(struct { @@ -93,7 +117,7 @@ func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType st Tools []map[string]any `json:"tools,omitempty"` Input []core.ResponsesInputElement `json:"input"` }{req.Instructions, req.Tools, items[:len(items)-1]}) - if err != nil || estimatedTokens(prefixBody) < providerCacheMinimum(providerType, selector.Model) { + if err != nil || estimatedTokens(prefixBody) < providerCacheMinimum(profile, selector.Model) { return req } planned, ok := cloneResponsesRequest(req) @@ -101,8 +125,8 @@ func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType st return req } key := cacheAffinityKey(providerType, selector, req.User, prefixBody) - switch normalizedProviderType(providerType) { - case "openai": + switch profile.mode { + case promptCacheOpenAI: planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ "prompt_cache_key": jsonString(key), }) @@ -111,7 +135,7 @@ func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType st "prompt_cache_options": json.RawMessage(`{"mode":"explicit"}`), }) } - case "anthropic": + case promptCacheAnthropic: planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ "cache_control": json.RawMessage(`{"type":"ephemeral"}`), }) @@ -128,6 +152,10 @@ func cloneChatRequest(req *core.ChatRequest) (*core.ChatRequest, bool) { if err := json.Unmarshal(body, &clone); err != nil { return nil, false } + if req.PromptCachePlan != nil { + plan := *req.PromptCachePlan + clone.PromptCachePlan = &plan + } return &clone, true } @@ -144,7 +172,7 @@ func cloneResponsesRequest(req *core.ResponsesRequest) (*core.ResponsesRequest, } func hasCacheDirective(fields core.UnknownJSONFields) bool { - for _, key := range []string{"cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options"} { + for _, key := range []string{"cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options", "prompt_cache_breakpoint", core.GatewayCachePointField} { if len(fields.Lookup(key)) > 0 { return true } @@ -152,9 +180,120 @@ func hasCacheDirective(fields core.UnknownJSONFields) bool { return false } +func hasChatCacheDirective(req *core.ChatRequest, prefix []core.Message) bool { + if hasCacheDirective(req.ExtraFields) || anyHasCacheDirective(req.Tools) { + return true + } + for _, message := range prefix { + if hasCacheDirective(message.ExtraFields) { + return true + } + for _, call := range message.ToolCalls { + if hasCacheDirective(call.ExtraFields) || hasCacheDirective(call.Function.ExtraFields) { + return true + } + } + if contentHasCacheDirective(message.Content) { + return true + } + } + return false +} + +func hasResponsesCacheDirective(req *core.ResponsesRequest, prefix []core.ResponsesInputElement) bool { + if hasCacheDirective(req.ExtraFields) || anyHasCacheDirective(req.Tools) { + return true + } + for _, item := range prefix { + if hasCacheDirective(item.ExtraFields) || contentHasCacheDirective(item.Content) { + return true + } + } + return false +} + +func contentHasCacheDirective(content any) bool { + switch typed := content.(type) { + case []core.ContentPart: + for _, part := range typed { + if hasCacheDirective(part.ExtraFields) || + (part.ImageURL != nil && hasCacheDirective(part.ImageURL.ExtraFields)) || + (part.InputAudio != nil && hasCacheDirective(part.InputAudio.ExtraFields)) { + return true + } + } + return false + default: + return anyHasCacheDirective(typed) + } +} + +func anyHasCacheDirective(value any) bool { + switch typed := value.(type) { + case map[string]any: + for key, child := range typed { + if isCacheDirectiveKey(key) || anyHasCacheDirective(child) { + return true + } + } + case []any: + if slices.ContainsFunc(typed, anyHasCacheDirective) { + return true + } + case []map[string]any: + for _, child := range typed { + if anyHasCacheDirective(child) { + return true + } + } + } + return false +} + +func isCacheDirectiveKey(key string) bool { + switch key { + case "cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options", "prompt_cache_breakpoint", core.GatewayCachePointField: + return true + default: + return false + } +} + +func estimateSimpleChatPrefixTokens(req *core.ChatRequest) (int, bool) { + if len(req.Tools) > 0 { + return 0, false + } + bytes := 0 + for _, message := range req.Messages[:len(req.Messages)-1] { + bytes += len(message.Role) + len(message.ToolCallID) + 16 + switch content := message.Content.(type) { + case string: + bytes += len(content) + case []core.ContentPart: + for _, part := range content { + if part.Type != "text" && part.Type != "input_text" { + return 0, false + } + bytes += len(part.Text) + 16 + } + default: + return 0, false + } + if len(message.ToolCalls) > 0 { + return 0, false + } + } + return (bytes + 3) / 4, true +} + func markLastChatPrefix(messages *[]core.Message, field string, value json.RawMessage) { for i := len(*messages) - 2; i >= 0; i-- { msg := &(*messages)[i] + switch msg.Role { + case "system", "developer", "user", "assistant": + default: + continue + } if strings.TrimSpace(core.ExtractTextContent(msg.Content)) == "" && len(msg.ToolCalls) == 0 { continue } @@ -214,7 +353,7 @@ func markOpenAIResponsesBreakpoint(req *core.ResponsesRequest) bool { return true case []core.ContentPart: for j := len(content) - 1; j >= 0; j-- { - if content[j].Type == "input_text" || content[j].Type == "input_image" || content[j].Type == "input_file" { + if content[j].Type == "text" || content[j].Type == "input_text" || content[j].Type == "input_image" || content[j].Type == "input_file" { content[j].ExtraFields = mergeCacheExtras(content[j].ExtraFields, map[string]json.RawMessage{ "prompt_cache_breakpoint": json.RawMessage(`{"mode":"explicit"}`), }) @@ -230,7 +369,7 @@ func markOpenAIResponsesBreakpoint(req *core.ResponsesRequest) bool { continue } blockType, _ := block["type"].(string) - if blockType == "input_text" || blockType == "input_image" || blockType == "input_file" { + if blockType == "text" || blockType == "input_text" || blockType == "input_image" || blockType == "input_file" { if _, exists := block["prompt_cache_breakpoint"]; !exists { block["prompt_cache_breakpoint"] = map[string]any{"mode": "explicit"} } @@ -239,6 +378,18 @@ func markOpenAIResponsesBreakpoint(req *core.ResponsesRequest) bool { return true } } + case []map[string]any: + for j := len(content) - 1; j >= 0; j-- { + blockType, _ := content[j]["type"].(string) + if blockType == "text" || blockType == "input_text" || blockType == "input_image" || blockType == "input_file" { + if _, exists := content[j]["prompt_cache_breakpoint"]; !exists { + content[j]["prompt_cache_breakpoint"] = map[string]any{"mode": "explicit"} + } + items[i].Content = content + req.Input = items + return true + } + } } } return false @@ -274,13 +425,12 @@ func cacheAffinityKey(providerType string, selector core.ModelSelector, user str func estimatedTokens(body []byte) int { return (len(body) + 3) / 4 } -func providerCacheMinimum(providerType, model string) int { - providerType = normalizedProviderType(providerType) +func providerCacheMinimum(profile promptCacheProfile, model string) int { model = strings.ToLower(model) - switch providerType { - case "openai": + switch profile.mode { + case promptCacheOpenAI: return 1024 - case "anthropic": + case promptCacheAnthropic: if strings.Contains(model, "haiku-3") && !strings.Contains(model, "3-5") && !strings.Contains(model, "3.5") { return 4096 } @@ -288,15 +438,38 @@ func providerCacheMinimum(providerType, model string) int { return 2048 } return 1024 - case "gemini", "vertex": + case promptCacheGemini: return 4096 - case "bedrock": - return 1024 + case promptCacheBedrock: + if strings.Contains(model, "nova") { + return 1536 + } + if strings.Contains(model, "claude") { + return 1024 + } + return int(^uint(0) >> 1) default: return int(^uint(0) >> 1) } } +func promptCacheProfileFor(providerType string) promptCacheProfile { + switch normalizedProviderType(providerType) { + case "openai": + return promptCacheProfile{mode: promptCacheOpenAI} + case "anthropic": + return promptCacheProfile{mode: promptCacheAnthropic, acceptsAnthropicCacheControl: true} + case "openrouter": + return promptCacheProfile{acceptsAnthropicCacheControl: true} + case "bedrock": + return promptCacheProfile{mode: promptCacheBedrock} + case "gemini": + return promptCacheProfile{mode: promptCacheGemini} + default: + return promptCacheProfile{} + } +} + func normalizedProviderType(providerType string) string { return strings.ToLower(strings.TrimSpace(providerType)) } diff --git a/internal/providers/cache_planner_test.go b/internal/providers/cache_planner_test.go index bb0d8050..24b651ce 100644 --- a/internal/providers/cache_planner_test.go +++ b/internal/providers/cache_planner_test.go @@ -1,6 +1,8 @@ package providers import ( + "bytes" + "os" "strings" "testing" @@ -10,6 +12,7 @@ import ( ) func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *testing.T) { + planner := &cachePlanner{enabled: true} prefix := strings.Repeat("stable context ", 1500) tests := []struct { provider string @@ -19,7 +22,7 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes }{ {provider: "openai", model: "gpt-5.6", field: "prompt_cache_key"}, {provider: "anthropic", model: "claude-sonnet-4-5", field: "cache_control"}, - {provider: "bedrock", model: "anthropic.claude-sonnet-4-5", marker: cachePointField}, + {provider: "bedrock", model: "anthropic.claude-sonnet-4-5", marker: core.GatewayCachePointField}, {provider: "gemini", model: "gemini-2.5-pro"}, } for _, tt := range tests { @@ -28,7 +31,7 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes {Role: "system", Content: prefix}, {Role: "user", Content: "new turn"}, }} - planned := newCachePlanner().planChat(req, tt.provider, core.ModelSelector{Provider: tt.provider + "-primary", Model: tt.model}) + planned := planner.planChat(req, tt.provider, core.ModelSelector{Provider: tt.provider + "-primary", Model: tt.model}) if planned == req { t.Fatal("planner returned caller-owned request") } @@ -58,16 +61,30 @@ func TestNewCachePlanner_EnvironmentKillSwitch(t *testing.T) { for _, tt := range []struct { name string value string + set bool enabled bool }{ {name: "default on", enabled: true}, - {name: "explicit on", value: "true", enabled: true}, - {name: "explicit off", value: "false", enabled: false}, - {name: "invalid keeps safe default", value: "sometimes", enabled: true}, + {name: "empty keeps safe default", value: "", set: true, enabled: true}, + {name: "explicit on", value: "true", set: true, enabled: true}, + {name: "explicit off", value: "false", set: true, enabled: false}, + {name: "invalid keeps safe default", value: "sometimes", set: true, enabled: true}, } { t.Run(tt.name, func(t *testing.T) { - t.Setenv(providerPromptCachePlannerEnabledEnv, tt.value) - if got := newCachePlanner() != nil; got != tt.enabled { + old, existed := os.LookupEnv(providerPromptCachePlannerEnabledEnv) + t.Cleanup(func() { + if existed { + _ = os.Setenv(providerPromptCachePlannerEnabledEnv, old) + } else { + _ = os.Unsetenv(providerPromptCachePlannerEnabledEnv) + } + }) + if tt.set { + _ = os.Setenv(providerPromptCachePlannerEnabledEnv, tt.value) + } else { + _ = os.Unsetenv(providerPromptCachePlannerEnabledEnv) + } + if got := newCachePlanner().enabled; got != tt.enabled { t.Fatalf("newCachePlanner() enabled = %v, want %v", got, tt.enabled) } }) @@ -75,7 +92,7 @@ func TestNewCachePlanner_EnvironmentKillSwitch(t *testing.T) { } func TestCachePlannerHonorsMinimumAndClientDirective(t *testing.T) { - planner := newCachePlanner() + planner := &cachePlanner{enabled: true} short := &core.ChatRequest{Messages: []core.Message{{Role: "system", Content: "short"}, {Role: "user", Content: "turn"}}} if got := planner.planChat(short, "openai", core.ModelSelector{Model: "gpt-5.6"}); got != short { t.Fatal("planned prefix below provider minimum") @@ -89,3 +106,96 @@ func TestCachePlannerHonorsMinimumAndClientDirective(t *testing.T) { t.Fatal("overrode client cache directive") } } + +func TestCachePlannerResponsesShapesAndCallerOwnership(t *testing.T) { + planner := &cachePlanner{enabled: true} + prefix := strings.Repeat("stable response context ", 1200) + shapes := []struct { + name string + content any + }{ + {name: "string", content: prefix}, + {name: "typed parts", content: []core.ContentPart{{Type: "input_text", Text: prefix}}}, + {name: "generic parts", content: []any{map[string]any{"type": "input_text", "text": prefix}}}, + } + for _, shape := range shapes { + t.Run(shape.name, func(t *testing.T) { + req := &core.ResponsesRequest{Model: "gpt-5.6", Input: []core.ResponsesInputElement{ + {Role: "user", Content: shape.content}, + {Role: "user", Content: "dynamic"}, + }} + before, err := json.Marshal(req) + if err != nil { + t.Fatal(err) + } + planned := planner.planResponses(req, "openai", core.ModelSelector{Provider: "openai-primary", Model: "gpt-5.6"}) + if planned == req || len(planned.ExtraFields.Lookup("prompt_cache_key")) == 0 || + len(planned.ExtraFields.Lookup("prompt_cache_options")) == 0 { + t.Fatalf("Responses plan missing cache fields: %+v", planned) + } + plannedJSON, err := json.Marshal(planned) + if err != nil || !bytes.Contains(plannedJSON, []byte(`"prompt_cache_breakpoint"`)) { + t.Fatalf("Responses plan lacks explicit breakpoint: %s (err=%v)", plannedJSON, err) + } + after, err := json.Marshal(req) + if err != nil { + t.Fatal(err) + } + if string(before) != string(after) { + t.Fatalf("planner mutated caller: before=%s after=%s", before, after) + } + }) + } + anthropic := &core.ResponsesRequest{Input: []core.ResponsesInputElement{ + {Role: "user", Content: prefix}, {Role: "user", Content: "dynamic"}, + }} + if got := planner.planResponses(anthropic, "anthropic", core.ModelSelector{Model: "claude-sonnet-4-5"}); got == anthropic || len(got.ExtraFields.Lookup("cache_control")) == 0 { + t.Fatal("Anthropic Responses plan lacks cache_control") + } + short := &core.ResponsesRequest{Input: []core.ResponsesInputElement{ + {Role: "user", Content: "short"}, {Role: "user", Content: "dynamic"}, + }} + if got := planner.planResponses(short, "openai", core.ModelSelector{Model: "gpt-5.6"}); got != short { + t.Fatal("planned a Responses prefix below the provider minimum") + } +} + +func TestCachePlannerFindsNestedClientDirective(t *testing.T) { + prefix := strings.Repeat("x", 9000) + req := &core.ChatRequest{Messages: []core.Message{ + {Role: "system", Content: []core.ContentPart{{ + Type: "text", Text: prefix, + ExtraFields: core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{ + "cache_control": json.RawMessage(`{"type":"ephemeral"}`), + }), + }}}, + {Role: "user", Content: "turn"}, + }} + if got := (&cachePlanner{enabled: true}).planChat(req, "openai", core.ModelSelector{Model: "gpt-5.6"}); got != req { + t.Fatal("planner overrode a nested client cache directive") + } +} + +func TestCachePlannerProviderCapabilityBoundaries(t *testing.T) { + req := &core.ChatRequest{Messages: []core.Message{ + {Role: "system", Content: strings.Repeat("x", 20000)}, + {Role: "user", Content: "turn"}, + }} + planner := &cachePlanner{enabled: true} + for _, provider := range []string{"openrouter", "vertex", "unknown"} { + if got := planner.planChat(req, provider, core.ModelSelector{Model: "gemini-2.5-pro"}); got != req { + t.Fatalf("provider %q unexpectedly received an automatic plan", provider) + } + } +} + +func TestCloneChatRequestPreservesInternalCachePlan(t *testing.T) { + req := &core.ChatRequest{PromptCachePlan: &core.PromptCachePlan{Key: "stable"}} + clone, ok := cloneChatRequest(req) + if !ok || clone.PromptCachePlan == nil || clone.PromptCachePlan.Key != "stable" { + t.Fatalf("clone lost internal cache metadata: %+v", clone) + } + if clone.PromptCachePlan == req.PromptCachePlan { + t.Fatal("clone aliases internal cache metadata") + } +} diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 9474c53c..ba4c8118 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -3,8 +3,11 @@ package gemini import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "io" + "log/slog" "net/http" "net/url" "os" @@ -63,6 +66,11 @@ const ( // accepts further spellings of each. geminiAPIModeNative = "native" geminiAPIModeOpenAICompatible = "openai_compatible" + geminiCacheTTL = 5 * time.Minute + geminiCacheFreshness = 15 * time.Second + geminiCacheFailureBackoff = 30 * time.Second + geminiCacheCreateTimeout = 30 * time.Second + geminiCacheObjectLimit = 1024 ) // Provider implements the core.Provider interface for Google Gemini @@ -82,8 +90,9 @@ type Provider struct { } type geminiCacheObject struct { - name string - expiresAt time.Time + name string + expiresAt time.Time + retryAfter time.Time } // New creates a new Gemini provider. @@ -550,22 +559,37 @@ type geminiCreateCachedContentResponse struct { // post-routing planner. Cache creation is best-effort: an unsupported model or // endpoint must never turn a request that would otherwise work into a failure. func (p *Provider) prepareCachedContent(ctx context.Context, req *core.ChatRequest, body *geminiGenerateContentRequest) { - if body == nil || body.CachedContent != "" || len(body.Contents) < 2 { + if body == nil || body.CachedContent != "" || p.backend != geminiBackendAIStudio || len(body.Contents) == 0 { return } if req.PromptCachePlan == nil || strings.TrimSpace(req.PromptCachePlan.Key) == "" { return } - key := req.PromptCachePlan.Key - if cached := p.cachedContentObject(key); cached != "" { - useGeminiCachedPrefix(body, cached) + // The final content is the live turn. A system instruction, tools, or an + // earlier content item must remain before it for a useful cached prefix. + if len(body.Contents) == 1 && body.SystemInstruction == nil && len(body.Tools) == 0 { + return + } + key, ok := p.scopedCachedContentKey(ctx, req.PromptCachePlan.Key) + if !ok { + return + } + if cached, suppress := p.cachedContentObject(key, time.Now()); suppress { + if cached != "" { + useGeminiCachedPrefix(body, cached) + } return } - value, err, _ := p.cacheFlight.Do(key, func() (any, error) { - if cached := p.cachedContentObject(key); cached != "" { + value, _, _ := p.cacheFlight.Do(key, func() (any, error) { + now := time.Now() + if cached, suppress := p.cachedContentObject(key, now); suppress { return cached, nil } + // Cache creation may outlive the request that happened to lead the + // singleflight. Keep affinity values while detaching cancellation. + createCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), geminiCacheCreateTimeout) + defer cancel() createReq := geminiCreateCachedContentRequest{ Model: "models/" + normalizeGeminiModelID(req.Model), SystemInstruction: body.SystemInstruction, @@ -574,47 +598,82 @@ func (p *Provider) prepareCachedContent(ctx context.Context, req *core.ChatReque TTL: "300s", } var created geminiCreateCachedContentResponse - if err := p.nativeClient.Do(ctx, llmclient.Request{ - Method: http.MethodPost, - Endpoint: "/cachedContents", - Body: &createReq, - }, &created); err != nil { - return "", err - } - if strings.TrimSpace(created.Name) == "" { - return "", fmt.Errorf("cached-content creation for Gemini returned an empty name") + err := p.nativeClient.Do(createCtx, llmclient.Request{ + Method: http.MethodPost, Endpoint: "/cachedContents", Body: &createReq, + }, &created) + if err != nil || strings.TrimSpace(created.Name) == "" { + if err == nil { + err = fmt.Errorf("cached-content creation returned an empty name") + } + slog.DebugContext(ctx, "Gemini cached-content creation skipped", "model", req.Model, "error", err) + p.storeCachedContentObject(key, geminiCacheObject{retryAfter: now.Add(geminiCacheFailureBackoff)}, now) + return "", nil } expiresAt := created.ExpireTime if expiresAt.IsZero() { - expiresAt = time.Now().Add(5 * time.Minute) + expiresAt = now.Add(geminiCacheTTL) } - p.cacheMu.Lock() - if p.cacheObjects == nil { - p.cacheObjects = make(map[string]geminiCacheObject) + if !now.Add(geminiCacheFreshness).Before(expiresAt) { + p.storeCachedContentObject(key, geminiCacheObject{retryAfter: now.Add(geminiCacheFailureBackoff)}, now) + return "", nil } - p.cacheObjects[key] = geminiCacheObject{name: created.Name, expiresAt: expiresAt} - p.cacheMu.Unlock() + p.storeCachedContentObject(key, geminiCacheObject{name: created.Name, expiresAt: expiresAt}, now) return created.Name, nil }) - if err == nil { - if cached, ok := value.(string); ok && cached != "" { - useGeminiCachedPrefix(body, cached) - } + if cached, ok := value.(string); ok && cached != "" { + useGeminiCachedPrefix(body, cached) + } +} + +func (p *Provider) scopedCachedContentKey(ctx context.Context, planKey string) (string, bool) { + credential, stable := p.keys.StableForContext(ctx) + if !stable { + return "", false } + digest := sha256.Sum256([]byte(credential)) + return planKey + ":" + hex.EncodeToString(digest[:]), true } -func (p *Provider) cachedContentObject(key string) string { +func (p *Provider) cachedContentObject(key string, now time.Time) (string, bool) { p.cacheMu.Lock() defer p.cacheMu.Unlock() entry, ok := p.cacheObjects[key] if !ok { - return "" + return "", false + } + if entry.name == "" { + if now.Before(entry.retryAfter) { + return "", true + } + delete(p.cacheObjects, key) + return "", false } - if time.Now().Add(15 * time.Second).After(entry.expiresAt) { + if now.Add(geminiCacheFreshness).After(entry.expiresAt) { delete(p.cacheObjects, key) - return "" + return "", false + } + return entry.name, true +} + +func (p *Provider) storeCachedContentObject(key string, entry geminiCacheObject, now time.Time) { + p.cacheMu.Lock() + defer p.cacheMu.Unlock() + if p.cacheObjects == nil { + p.cacheObjects = make(map[string]geminiCacheObject) + } + for candidate, cached := range p.cacheObjects { + if (cached.name == "" && !now.Before(cached.retryAfter)) || + (cached.name != "" && now.Add(geminiCacheFreshness).After(cached.expiresAt)) { + delete(p.cacheObjects, candidate) + } + } + if len(p.cacheObjects) >= geminiCacheObjectLimit { + for candidate := range p.cacheObjects { + delete(p.cacheObjects, candidate) + break + } } - return entry.name + p.cacheObjects[key] = entry } func useGeminiCachedPrefix(body *geminiGenerateContentRequest, name string) { diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 9826b721..436a41da 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -1,15 +1,19 @@ package gemini import ( + "bytes" "context" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" "strings" + "sync" "sync/atomic" "testing" + "time" "github.com/enterpilot/gomodel/internal/core" "github.com/enterpilot/gomodel/internal/llmclient" @@ -37,9 +41,12 @@ func TestNew(t *testing.T) { func TestPrepareCachedContentCreatesAndReusesObject(t *testing.T) { var creates atomic.Int32 + var wrongPath atomic.Bool server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/cachedContents" { - t.Fatalf("path = %q, want /cachedContents", r.URL.Path) + wrongPath.Store(true) + http.Error(w, "wrong path", http.StatusNotFound) + return } creates.Add(1) w.Header().Set("Content-Type", "application/json") @@ -57,6 +64,7 @@ func TestPrepareCachedContentCreatesAndReusesObject(t *testing.T) { {Role: "user", Parts: []geminiPart{{Text: "stable"}}}, {Role: "user", Parts: []geminiPart{{Text: "dynamic"}}}, }, + Tools: []geminiTool{{FunctionDeclarations: []geminiFunctionDeclaration{{Name: "lookup"}}}}, } } first := newBody() @@ -66,6 +74,9 @@ func TestPrepareCachedContentCreatesAndReusesObject(t *testing.T) { if creates.Load() != 1 { t.Fatalf("cache creates = %d, want 1", creates.Load()) } + if wrongPath.Load() { + t.Fatal("cached-content request used the wrong path") + } for i, body := range []*geminiGenerateContentRequest{first, second} { if body.CachedContent != "cachedContents/session-prefix" || len(body.Contents) != 1 || body.SystemInstruction != nil { t.Fatalf("body %d did not use cached prefix: %+v", i, body) @@ -73,6 +84,179 @@ func TestPrepareCachedContentCreatesAndReusesObject(t *testing.T) { } } +func TestPrepareCachedContentSupportsSystemOnlyPrefix(t *testing.T) { + var creates atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + creates.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"name":"cachedContents/system-prefix","expireTime":"2099-01-01T00:00:00Z"}`) + })) + defer server.Close() + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) + req := &core.ChatRequest{Model: "gemini-2.5-pro", PromptCachePlan: &core.PromptCachePlan{Key: "system"}} + body := &geminiGenerateContentRequest{ + SystemInstruction: &geminiContent{Parts: []geminiPart{{Text: "stable system"}}}, + Contents: []geminiContent{{Role: "user", Parts: []geminiPart{{Text: "live turn"}}}}, + } + p.prepareCachedContent(context.Background(), req, body) + if creates.Load() != 1 || body.CachedContent != "cachedContents/system-prefix" || len(body.Contents) != 1 || body.SystemInstruction != nil { + t.Fatalf("system prefix was not cached: creates=%d body=%+v", creates.Load(), body) + } +} + +func TestPrepareCachedContentFailureIsBestEffortAndBackedOff(t *testing.T) { + var creates atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + creates.Add(1) + http.Error(w, "unsupported", http.StatusBadRequest) + })) + defer server.Close() + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) + req := &core.ChatRequest{Model: "gemini-2.5-pro", PromptCachePlan: &core.PromptCachePlan{Key: "failure"}} + newBody := func() *geminiGenerateContentRequest { + return &geminiGenerateContentRequest{ + SystemInstruction: &geminiContent{Parts: []geminiPart{{Text: "system"}}}, + Contents: []geminiContent{{Role: "user", Parts: []geminiPart{{Text: "live"}}}}, + Tools: []geminiTool{{FunctionDeclarations: []geminiFunctionDeclaration{{Name: "lookup"}}}}, + } + } + first, second := newBody(), newBody() + firstBefore, _ := json.Marshal(first) + secondBefore, _ := json.Marshal(second) + p.prepareCachedContent(context.Background(), req, first) + p.prepareCachedContent(context.Background(), req, second) + if creates.Load() != 1 { + t.Fatalf("failed creation attempts = %d, want one during backoff", creates.Load()) + } + firstAfter, _ := json.Marshal(first) + secondAfter, _ := json.Marshal(second) + if !bytes.Equal(firstBefore, firstAfter) || !bytes.Equal(secondBefore, secondAfter) { + t.Fatalf("failed best-effort creation modified requests: %s / %s", firstAfter, secondAfter) + } +} + +func TestPrepareCachedContentEmptyNameAndExpiringEntry(t *testing.T) { + var creates atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + call := creates.Add(1) + w.Header().Set("Content-Type", "application/json") + if call == 1 { + _, _ = io.WriteString(w, `{}`) + return + } + _, _ = io.WriteString(w, `{"name":"cachedContents/recreated","expireTime":"2099-01-01T00:00:00Z"}`) + })) + defer server.Close() + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) + req := &core.ChatRequest{Model: "gemini-2.5-pro", PromptCachePlan: &core.PromptCachePlan{Key: "expiry"}} + newBody := func() *geminiGenerateContentRequest { + return &geminiGenerateContentRequest{SystemInstruction: &geminiContent{}, Contents: []geminiContent{{Role: "user"}}} + } + failed := newBody() + before, _ := json.Marshal(failed) + p.prepareCachedContent(context.Background(), req, failed) + after, _ := json.Marshal(failed) + if !bytes.Equal(before, after) { + t.Fatalf("empty-name failure modified request: %s", after) + } + scopedKey, ok := p.scopedCachedContentKey(context.Background(), req.PromptCachePlan.Key) + if !ok { + t.Fatal("single credential did not produce a stable key") + } + p.cacheMu.Lock() + p.cacheObjects[scopedKey] = geminiCacheObject{name: "cachedContents/expiring", expiresAt: time.Now().Add(5 * time.Second)} + p.cacheMu.Unlock() + recreated := newBody() + p.prepareCachedContent(context.Background(), req, recreated) + if creates.Load() != 2 || recreated.CachedContent != "cachedContents/recreated" { + t.Fatalf("expiring entry was not recreated: creates=%d body=%+v", creates.Load(), recreated) + } +} + +func TestPrepareCachedContentCoalescesConcurrentCreation(t *testing.T) { + const callers = 12 + var creates, begun atomic.Int32 + allBegun := make(chan struct{}) + handlerEntered := make(chan struct{}, 1) + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + creates.Add(1) + handlerEntered <- struct{}{} + <-allBegun + <-release + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"name":"cachedContents/concurrent","expireTime":"2099-01-01T00:00:00Z"}`) + })) + defer server.Close() + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) + req := &core.ChatRequest{Model: "gemini-2.5-pro", PromptCachePlan: &core.PromptCachePlan{Key: "concurrent"}} + bodies := make([]*geminiGenerateContentRequest, callers) + var wg sync.WaitGroup + for i := range callers { + wg.Go(func() { + bodies[i] = &geminiGenerateContentRequest{SystemInstruction: &geminiContent{}, Contents: []geminiContent{{Role: "user"}}} + if begun.Add(1) == callers { + close(allBegun) + } + p.prepareCachedContent(context.Background(), req, bodies[i]) + }) + } + <-handlerEntered + <-allBegun + close(release) + wg.Wait() + if got := creates.Load(); got != 1 { + t.Fatalf("concurrent cache creates = %d, want 1", got) + } + for i, body := range bodies { + if body.CachedContent != "cachedContents/concurrent" { + t.Fatalf("caller %d did not receive shared cache object: %+v", i, body) + } + } +} + +func TestPrepareCachedContentRequiresStableCredentialAndAIStudio(t *testing.T) { + p := NewWithHTTPClient("key", http.DefaultClient, llmclient.Hooks{}) + p.keys = providers.NewKeyring("one", "two") + req := &core.ChatRequest{PromptCachePlan: &core.PromptCachePlan{Key: "prefix"}} + body := &geminiGenerateContentRequest{SystemInstruction: &geminiContent{}, Contents: []geminiContent{{Role: "user"}}} + if _, ok := p.scopedCachedContentKey(context.Background(), "prefix"); ok { + t.Fatal("sessionless rotating credentials must not own reusable cache objects") + } + ctx := core.WithSessionID(context.Background(), "session-a") + first, ok := p.scopedCachedContentKey(ctx, "prefix") + second, ok2 := p.scopedCachedContentKey(ctx, "prefix") + if !ok || !ok2 || first == "" || first != second { + t.Fatalf("sticky credential key is unstable: %q %q", first, second) + } + p.backend = geminiBackendVertex + p.prepareCachedContent(ctx, req, body) + if body.CachedContent != "" { + t.Fatal("Vertex must not use the AI Studio cachedContents endpoint") + } +} + +func TestGeminiCacheObjectMapIsBoundedAndSweepsExpiredEntries(t *testing.T) { + p := &Provider{} + now := time.Now() + p.cacheObjects = map[string]geminiCacheObject{ + "expired": {name: "old", expiresAt: now.Add(-time.Minute)}, + } + for i := range geminiCacheObjectLimit + 20 { + p.storeCachedContentObject(fmt.Sprintf("key-%d", i), geminiCacheObject{name: "cache", expiresAt: now.Add(time.Hour)}, now) + } + if _, exists := p.cacheObjects["expired"]; exists { + t.Fatal("expired entry was not swept") + } + if got := len(p.cacheObjects); got > geminiCacheObjectLimit { + t.Fatalf("cache object map size = %d, limit = %d", got, geminiCacheObjectLimit) + } +} + func TestNew_ReturnsProvider(t *testing.T) { provider := New(providers.ProviderConfig{APIKey: "test-api-key"}, providers.ProviderOptions{}) diff --git a/internal/providers/keyring.go b/internal/providers/keyring.go index e4a4c30d..9bcc510a 100644 --- a/internal/providers/keyring.go +++ b/internal/providers/keyring.go @@ -85,6 +85,24 @@ func (k *Keyring) NextForContext(ctx context.Context) string { return k.NextForSession(core.SessionIDFromContext(ctx)) } +// StableForContext returns the credential that every request with this context +// will use without advancing round-robin state. The boolean is false when the +// context has no stable credential affinity, so callers must not persist +// credential-scoped provider resources for later reuse. +func (k *Keyring) StableForContext(ctx context.Context) (string, bool) { + if k == nil || len(k.keys) == 0 { + return "", false + } + if len(k.keys) == 1 { + return k.keys[0], true + } + sessionID := core.SessionIDFromContext(ctx) + if !k.sessionSticky || sessionID == "" { + return "", false + } + return k.stickyKey(sessionID), true +} + // NextForSession deterministically maps one non-empty session to one key using // rendezvous hashing. Adding or removing a key therefore remaps only sessions // assigned to the changed key, rather than invalidating every warm cache. The @@ -99,6 +117,10 @@ func (k *Keyring) NextForSession(sessionID string) string { if !k.sessionSticky || sessionID == "" { return k.Next() } + return k.stickyKey(sessionID) +} + +func (k *Keyring) stickyKey(sessionID string) string { selected := k.keys[0] best := rendezvousKeyScore(sessionID, selected) for _, key := range k.keys[1:] { diff --git a/internal/providers/keyring_test.go b/internal/providers/keyring_test.go index eace5936..b73f8790 100644 --- a/internal/providers/keyring_test.go +++ b/internal/providers/keyring_test.go @@ -82,6 +82,35 @@ func TestKeyringSessionlessTrafficRemainsRoundRobin(t *testing.T) { } } +func TestKeyringStableForContext(t *testing.T) { + sticky := NewKeyring("k1", "k2") + ctx := core.WithSessionID(context.Background(), "conversation-42") + want := sticky.NextForContext(ctx) + for range 3 { + if got, ok := sticky.StableForContext(ctx); !ok || got != want { + t.Fatalf("StableForContext() = %q, %v, want %q, true", got, ok, want) + } + } + if got := sticky.Next(); got != "k1" { + t.Fatalf("stable lookup advanced round robin to %q", got) + } + + for name, ring := range map[string]*Keyring{ + "sessionless": sticky, + "disabled": NewKeyringWithSessionStickiness(false, "k1", "k2"), + } { + t.Run(name, func(t *testing.T) { + testCtx := context.Background() + if name == "disabled" { + testCtx = ctx + } + if key, ok := ring.StableForContext(testCtx); ok || key != "" { + t.Fatalf("StableForContext() = %q, %v, want empty, false", key, ok) + } + }) + } +} + func TestKeyringRendezvousHashingOnlyRemapsSessionsOnChangedKey(t *testing.T) { before := NewKeyring("k1", "k2", "k3") after := NewKeyring("k1", "k3") diff --git a/internal/providers/router.go b/internal/providers/router.go index e5b814c0..eebba38c 100644 --- a/internal/providers/router.go +++ b/internal/providers/router.go @@ -582,7 +582,7 @@ func forwardResponsesRequest(req *core.ResponsesRequest, selector core.ModelSele func (r *Router) plannedChatRequest(ctx context.Context, req *core.ChatRequest, selector core.ModelSelector) *core.ChatRequest { forward := r.forwardChatRequest(ctx, req, selector) - if r.cachePlanner == nil || len(forward.Messages) < 2 { + if r.cachePlanner == nil { return forward } return r.cachePlanner.planChat(forward, r.lookup.GetProviderType(selector.QualifiedModel()), selector) @@ -593,9 +593,6 @@ func (r *Router) plannedResponsesRequest(req *core.ResponsesRequest, selector co if r.cachePlanner == nil { return forward } - if items, ok := forward.Input.([]core.ResponsesInputElement); !ok || len(items) < 2 { - return forward - } return r.cachePlanner.planResponses(forward, r.lookup.GetProviderType(selector.QualifiedModel()), selector) } diff --git a/internal/responsecache/exact_cache_test.go b/internal/responsecache/exact_cache_test.go index c4956559..b8057908 100644 --- a/internal/responsecache/exact_cache_test.go +++ b/internal/responsecache/exact_cache_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "sync" @@ -26,6 +27,38 @@ type concurrentTrackingStore struct { releaseCh chan struct{} } +type blockingMissExchange struct { + ctx context.Context + started chan<- struct{} + release <-chan struct{} + joined chan struct{} + contextCalls atomic.Int32 +} + +func (e *blockingMissExchange) Context() context.Context { + if e.joined != nil && e.contextCalls.Add(1) == 2 { + close(e.joined) + } + return e.ctx +} +func (e *blockingMissExchange) Path() string { return "/v1/chat/completions" } +func (e *blockingMissExchange) Method() string { return http.MethodPost } +func (e *blockingMissExchange) RequestHeader(string) string { return "" } +func (e *blockingMissExchange) MarkHit(string) {} +func (e *blockingMissExchange) ReplayHit([]byte, []byte, string) error { return nil } +func (e *blockingMissExchange) Capture(_ string, next func() error) ([]byte, bool, error) { + if e.started != nil { + e.started <- struct{}{} + } + if e.release != nil { + <-e.release + } + if err := next(); err != nil { + return nil, false, err + } + return []byte(`{"ok":true}`), true, nil +} + func newConcurrentTrackingStore() *concurrentTrackingStore { return &concurrentTrackingStore{ enterCh: make(chan struct{}, 1024), @@ -86,6 +119,20 @@ func driveHandleRequest( next func(c *echo.Context) error, ) *httptest.ResponseRecorder { t.Helper() + rec, err := driveHandleRequestResult(mw, workflow, body, headers, next) + if err != nil { + t.Fatalf("HandleRequest: %v", err) + } + return rec +} + +func driveHandleRequestResult( + mw *ResponseCacheMiddleware, + workflow *core.Workflow, + body []byte, + headers map[string]string, + next func(c *echo.Context) error, +) (*httptest.ResponseRecorder, error) { e := echo.New() req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") @@ -97,10 +144,8 @@ func driveHandleRequest( } rec := httptest.NewRecorder() c := e.NewContext(req, rec) - if err := mw.HandleRequest(c, body, func() error { return next(c) }); err != nil { - t.Fatalf("HandleRequest: %v", err) - } - return rec + err := mw.HandleRequest(c, body, func() error { return next(c) }) + return rec, err } func TestHandleRequest_ExactCacheHit(t *testing.T) { @@ -165,60 +210,142 @@ func TestHandleRequest_DifferentBodyDifferentKey(t *testing.T) { } } -func TestHandleRequest_CoalescesConcurrentIdenticalMisses(t *testing.T) { - store := cache.NewMapStore() - defer store.Close() - mw := NewResponseCacheMiddlewareWithStore(store, time.Hour) - workflow := resolvedWorkflow("openai", "gpt-4") +func TestStoreAfter_CoalescesConcurrentIdenticalMisses(t *testing.T) { + m := newSimpleCacheMiddleware(cache.NewMapStore(), time.Hour, nil) + defer m.close() body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) - var calls atomic.Int32 - started := make(chan struct{}) + started := make(chan struct{}, 1) release := make(chan struct{}) - next := func(c *echo.Context) error { - if calls.Add(1) == 1 { - close(started) - } - <-release - return c.JSON(http.StatusOK, map[string]string{"result": "shared"}) - } - const requests = 12 - recorders := make([]*httptest.ResponseRecorder, requests) + errs := make([]error, requests) var wg sync.WaitGroup - for i := range requests { + wg.Go(func() { + errs[0] = m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), started: started, release: release, + }, body, func() error { + calls.Add(1) + return nil + }) + }) + <-started + + joined := make([]chan struct{}, requests-1) + for i := 1; i < requests; i++ { + joined[i-1] = make(chan struct{}) wg.Go(func() { - recorders[i] = driveHandleRequest(t, mw, workflow, body, nil, next) + errs[i] = m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), joined: joined[i-1], + }, body, func() error { + calls.Add(1) + return nil + }) }) } - <-started - time.Sleep(20 * time.Millisecond) + for _, waiter := range joined { + <-waiter + } close(release) wg.Wait() if got := calls.Load(); got != 1 { t.Fatalf("provider calls = %d, want one coalesced miss", got) } - hits := 0 - for i, rec := range recorders { - if rec.Code != http.StatusOK || !bytes.Contains(rec.Body.Bytes(), []byte("shared")) { - t.Fatalf("response %d = status %d body %q", i, rec.Code, rec.Body.String()) - } - if rec.Header().Get("X-Cache") == "HIT (exact)" { - hits++ + for i := range requests { + if errs[i] != nil { + t.Fatalf("request %d: %v", i, errs[i]) } } - if hits != requests-1 { - t.Fatalf("coalesced hit responses = %d, want %d", hits, requests-1) +} + +func TestStoreAfter_CanceledFollowerDoesNotWaitForLeader(t *testing.T) { + store := cache.NewMapStore() + m := newSimpleCacheMiddleware(store, time.Hour, nil) + defer m.close() + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) + started := make(chan struct{}, 1) + release := make(chan struct{}) + leaderDone := make(chan error, 1) + go func() { + leaderDone <- m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), started: started, release: release, + }, body, func() error { return nil }) + }() + <-started + + followerCtx, cancel := context.WithCancel(context.Background()) + cancel() + follower := &blockingMissExchange{ctx: followerCtx} + if err := m.StoreAfter(follower, body, func() error { return nil }); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled follower error = %v, want context.Canceled", err) + } + close(release) + if err := <-leaderDone; err != nil { + t.Fatalf("leader error: %v", err) + } +} + +func TestStoreAfter_LeaderErrorIsNotFannedOut(t *testing.T) { + m := newSimpleCacheMiddleware(cache.NewMapStore(), time.Hour, nil) + defer m.close() + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) + started := make(chan struct{}, 1) + release := make(chan struct{}) + leaderErr := errors.New("first provider failed") + leaderDone := make(chan error, 1) + go func() { + leaderDone <- m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), started: started, release: release, + }, body, func() error { return leaderErr }) + }() + <-started + + joined := make(chan struct{}) + var followerCalls atomic.Int32 + followerDone := make(chan error, 1) + go func() { + followerDone <- m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), joined: joined, + }, body, func() error { + followerCalls.Add(1) + return nil + }) + }() + <-joined + close(release) + if err := <-leaderDone; !errors.Is(err, leaderErr) { + t.Fatalf("leader error = %v, want %v", err, leaderErr) + } + if err := <-followerDone; err != nil { + t.Fatalf("follower inherited leader error: %v", err) + } + if got := followerCalls.Load(); got != 1 { + t.Fatalf("follower provider calls = %d, want independent retry", got) } } func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { plan := resolvedWorkflow("openai", "gpt-4") - compact := []byte(`{"input":[1,2],"model":"gpt-4"}`) - formatted := []byte("{\n \"model\": \"gpt-4\",\n \"input\": [1, 2]\n}") - if first, second := hashRequest("/v1/embeddings", compact, plan), hashRequest("/v1/embeddings", formatted, plan); first != second { - t.Fatalf("canonical-equivalent JSON produced different keys: %s != %s", first, second) + for _, tt := range []struct { + name string + first string + second string + equal bool + }{ + {name: "formatting and key order", first: `{"input":[1,2],"model":"gpt-4"}`, second: "{\n \"model\": \"gpt-4\",\n \"input\": [1, 2]\n}", equal: true}, + {name: "nested key order", first: `{"input":{"b":2,"a":1},"model":"gpt-4"}`, second: `{"model":"gpt-4","input":{"a":1,"b":2}}`, equal: true}, + {name: "number spelling preserved", first: `{"input":1,"model":"gpt-4"}`, second: `{"input":1.0,"model":"gpt-4"}`, equal: false}, + {name: "array order matters", first: `{"input":[1,2],"model":"gpt-4"}`, second: `{"input":[2,1],"model":"gpt-4"}`, equal: false}, + {name: "malformed falls back exactly", first: `{"input":`, second: ` {"input":`, equal: false}, + {name: "multiple values fall back exactly", first: `{"a":1} {"b":2}`, second: `{"a":1} {"b":2}`, equal: false}, + } { + t.Run(tt.name, func(t *testing.T) { + first := hashRequest("/v1/embeddings", []byte(tt.first), plan) + second := hashRequest("/v1/embeddings", []byte(tt.second), plan) + if got := first == second; got != tt.equal { + t.Fatalf("key equality = %v, want %v: %s / %s", got, tt.equal, first, second) + } + }) } } diff --git a/internal/responsecache/simple.go b/internal/responsecache/simple.go index c23aad13..92efc89b 100644 --- a/internal/responsecache/simple.go +++ b/internal/responsecache/simple.go @@ -15,7 +15,6 @@ import ( "github.com/labstack/echo/v5" "github.com/tidwall/gjson" - "golang.org/x/sync/singleflight" "github.com/enterpilot/gomodel/internal/cache" "github.com/enterpilot/gomodel/internal/core" @@ -48,7 +47,14 @@ type simpleCacheMiddleware struct { workers sync.WaitGroup mu sync.RWMutex closed bool - misses singleflight.Group + missMu sync.Mutex + misses map[string]*exactMissCall +} + +type exactMissCall struct { + done chan struct{} + data []byte + err error } func newSimpleCacheMiddleware(store cache.Store, ttl time.Duration, hitRecorder func(exchange, []byte, string)) *simpleCacheMiddleware { @@ -103,45 +109,66 @@ func (m *simpleCacheMiddleware) StoreAfter(ex exchange, body []byte, next func() plan := core.GetWorkflow(ex.Context()) key := hashRequest(path, body, plan) - type missResult struct { - owner *struct{ marker byte } - data []byte - } - owner := &struct{ marker byte }{} - value, err, _ := m.misses.Do(key, func() (any, error) { + call, leader := m.joinMiss(key) + if leader { data, ok, err := ex.Capture("response cache: failed to capture cacheable response body", next) - if err != nil { - return nil, err + if err == nil && ok { + m.enqueueWrite(cacheWriteJob{key: key, data: data}) + } else { + data = nil } - if !ok { - return missResult{owner: owner}, nil - } - m.enqueueWrite(cacheWriteJob{key: key, data: data}) - return missResult{owner: owner, data: data}, nil - }) - if err != nil { + m.finishMiss(key, call, data, err) return err } - result, _ := value.(missResult) - if result.owner == owner { - return nil + + select { + case <-ex.Context().Done(): + return ex.Context().Err() + case <-call.done: } // The leader produced a non-cacheable result (failure status, failover, or // malformed body). Waiting followers must execute independently rather than // replaying something the normal cache would refuse to store. - if len(result.data) == 0 { + if call.err != nil { + return m.StoreAfter(ex, body, next) + } + if len(call.data) == 0 { return next() } - if err := ex.ReplayHit(body, result.data, CacheTypeExact); err != nil { + if err := ex.ReplayHit(body, call.data, CacheTypeExact); err != nil { return next() } ex.MarkHit(CacheTypeExact) if m.hitRecorder != nil { - m.hitRecorder(ex, result.data, CacheTypeExact) + m.hitRecorder(ex, call.data, CacheTypeExact) } return nil } +func (m *simpleCacheMiddleware) joinMiss(key string) (*exactMissCall, bool) { + m.missMu.Lock() + defer m.missMu.Unlock() + if call := m.misses[key]; call != nil { + return call, false + } + if m.misses == nil { + m.misses = make(map[string]*exactMissCall) + } + call := &exactMissCall{done: make(chan struct{})} + m.misses[key] = call + return call, true +} + +func (m *simpleCacheMiddleware) finishMiss(key string, call *exactMissCall, data []byte, err error) { + m.missMu.Lock() + call.data, call.err = data, err + if m.misses[key] == call { + delete(m.misses, key) + } + close(call.done) + m.missMu.Unlock() +} + // close waits for all in-flight cache writes to complete, then closes the store. func (m *simpleCacheMiddleware) close() error { m.mu.Lock() From 9829f0e5b2bc87a0a8732ebd53f4421208d6253b Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sat, 1 Aug 2026 21:30:14 +0200 Subject: [PATCH 4/6] fix(cache): address planner review follow-ups --- docs/pro.mdx | 4 ++ internal/providers/cache_planner.go | 18 ++++++- internal/providers/cache_planner_test.go | 61 +++++++++++++++++++++- internal/providers/gemini/gemini.go | 3 +- internal/responsecache/exact_cache_test.go | 3 ++ 5 files changed, 85 insertions(+), 4 deletions(-) diff --git a/docs/pro.mdx b/docs/pro.mdx index 7bd85446..991b3fde 100644 --- a/docs/pro.mdx +++ b/docs/pro.mdx @@ -110,6 +110,10 @@ prompt-cache reuse. ## Main settings +These variables are consumed by the separately distributed GoModel Pro +compression engine. They are intentionally not parsed by the open-core binary +when the Pro extension is absent. + | Env var | Default | Meaning | |---|---|---| | `PRO_COMPRESSION_ENABLED` | `true` | Master switch | diff --git a/internal/providers/cache_planner.go b/internal/providers/cache_planner.go index dafa1723..9ee3e58f 100644 --- a/internal/providers/cache_planner.go +++ b/internal/providers/cache_planner.go @@ -426,12 +426,18 @@ func cacheAffinityKey(providerType string, selector core.ModelSelector, user str func estimatedTokens(body []byte) int { return (len(body) + 3) / 4 } func providerCacheMinimum(profile promptCacheProfile, model string) int { - model = strings.ToLower(model) + model = strings.NewReplacer(".", "-", "_", "-").Replace(strings.ToLower(model)) switch profile.mode { case promptCacheOpenAI: return 1024 case promptCacheAnthropic: - if strings.Contains(model, "haiku-3") && !strings.Contains(model, "3-5") && !strings.Contains(model, "3.5") { + if strings.Contains(model, "fable-5") || strings.Contains(model, "mythos-5") { + return 512 + } + if strings.Contains(model, "mythos-preview") || strings.Contains(model, "opus-4-7") { + return 2048 + } + if strings.Contains(model, "haiku-4-5") || strings.Contains(model, "opus-4-5") || strings.Contains(model, "opus-4-6") { return 4096 } if strings.Contains(model, "haiku") { @@ -445,6 +451,10 @@ func providerCacheMinimum(profile promptCacheProfile, model string) int { return 1536 } if strings.Contains(model, "claude") { + if strings.Contains(model, "haiku-4-5") || strings.Contains(model, "sonnet-4-5") || + strings.Contains(model, "opus-4-5") || strings.Contains(model, "opus-4-6") { + return 4096 + } return 1024 } return int(^uint(0) >> 1) @@ -463,6 +473,10 @@ func promptCacheProfileFor(providerType string) promptCacheProfile { return promptCacheProfile{acceptsAnthropicCacheControl: true} case "bedrock": return promptCacheProfile{mode: promptCacheBedrock} + case "bedrock-mantle": + // Mantle exposes Bedrock models through the OpenAI-compatible cache + // fields, not Converse cachePoint blocks. + return promptCacheProfile{mode: promptCacheOpenAI} case "gemini": return promptCacheProfile{mode: promptCacheGemini} default: diff --git a/internal/providers/cache_planner_test.go b/internal/providers/cache_planner_test.go index 24b651ce..a51d44dd 100644 --- a/internal/providers/cache_planner_test.go +++ b/internal/providers/cache_planner_test.go @@ -23,6 +23,7 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes {provider: "openai", model: "gpt-5.6", field: "prompt_cache_key"}, {provider: "anthropic", model: "claude-sonnet-4-5", field: "cache_control"}, {provider: "bedrock", model: "anthropic.claude-sonnet-4-5", marker: core.GatewayCachePointField}, + {provider: "bedrock-mantle", model: "gpt-5.6", field: "prompt_cache_key"}, {provider: "gemini", model: "gemini-2.5-pro"}, } for _, tt := range tests { @@ -44,7 +45,7 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes if tt.provider == "gemini" && (planned.PromptCachePlan == nil || planned.PromptCachePlan.Key == "") { t.Fatal("Gemini plan lacks an internal cached-content key") } - if tt.provider == "openai" { + if promptCacheProfileFor(tt.provider).mode == promptCacheOpenAI { parts, ok := planned.Messages[0].Content.([]core.ContentPart) if !ok || len(parts) != 1 || len(parts[0].ExtraFields.Lookup("prompt_cache_breakpoint")) == 0 { t.Fatalf("OpenAI stable content lacks a breakpoint: %#v", planned.Messages[0].Content) @@ -57,6 +58,33 @@ func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *tes } } +func TestProviderCacheMinimumByModelGeneration(t *testing.T) { + for _, tt := range []struct { + provider string + model string + want int + }{ + {provider: "anthropic", model: "claude-haiku-4-5-20251001", want: 4096}, + {provider: "anthropic", model: "claude-haiku-3-5-latest", want: 2048}, + {provider: "anthropic", model: "claude-3-haiku-20240307", want: 2048}, + {provider: "anthropic", model: "claude-opus-4-5-20251101", want: 4096}, + {provider: "anthropic", model: "claude-opus-4.6", want: 4096}, + {provider: "anthropic", model: "claude-opus-4-7", want: 2048}, + {provider: "anthropic", model: "claude-sonnet-4-5", want: 1024}, + {provider: "bedrock", model: "anthropic.claude-haiku-4-5-20251001-v1:0", want: 4096}, + {provider: "bedrock", model: "anthropic.claude-sonnet-4-5-20250929-v1:0", want: 4096}, + {provider: "bedrock", model: "anthropic.claude-sonnet-4-6", want: 1024}, + {provider: "bedrock-mantle", model: "openai.gpt-5.6-sol", want: 1024}, + } { + t.Run(tt.provider+"/"+tt.model, func(t *testing.T) { + profile := promptCacheProfileFor(tt.provider) + if got := providerCacheMinimum(profile, tt.model); got != tt.want { + t.Fatalf("providerCacheMinimum() = %d, want %d", got, tt.want) + } + }) + } +} + func TestNewCachePlanner_EnvironmentKillSwitch(t *testing.T) { for _, tt := range []struct { name string @@ -199,3 +227,34 @@ func TestCloneChatRequestPreservesInternalCachePlan(t *testing.T) { t.Fatal("clone aliases internal cache metadata") } } + +func TestGeminiPlanKeyIncludesEntireNativePrefixAndBoundary(t *testing.T) { + planner := &cachePlanner{enabled: true} + makeRequest := func(system, boundary, toolName string) *core.ChatRequest { + return &core.ChatRequest{ + Model: "gemini-2.5-pro", + Messages: []core.Message{ + {Role: "system", Content: strings.Repeat(system, 5000)}, + {Role: "user", Content: boundary}, + {Role: "user", Content: "live"}, + }, + Tools: []map[string]any{{"type": "function", "function": map[string]any{"name": toolName}}}, + } + } + keys := make(map[string]struct{}) + for _, req := range []*core.ChatRequest{ + makeRequest("system-a", "boundary-a", "lookup"), + makeRequest("system-b", "boundary-a", "lookup"), + makeRequest("system-a", "boundary-b", "lookup"), + makeRequest("system-a", "boundary-a", "search"), + } { + planned := planner.planChat(req, "gemini", core.ModelSelector{Provider: "gemini-primary", Model: req.Model}) + if planned.PromptCachePlan == nil || planned.PromptCachePlan.Key == "" { + t.Fatal("Gemini request was not planned") + } + keys[planned.PromptCachePlan.Key] = struct{}{} + } + if len(keys) != 4 { + t.Fatalf("system, boundary, or tools were omitted from Gemini keys: %v", keys) + } +} diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index ba4c8118..a6764518 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -595,12 +595,13 @@ func (p *Provider) prepareCachedContent(ctx context.Context, req *core.ChatReque SystemInstruction: body.SystemInstruction, Contents: append([]geminiContent(nil), body.Contents[:len(body.Contents)-1]...), Tools: append([]geminiTool(nil), body.Tools...), - TTL: "300s", + TTL: fmt.Sprintf("%ds", geminiCacheTTL/time.Second), } var created geminiCreateCachedContentResponse err := p.nativeClient.Do(createCtx, llmclient.Request{ Method: http.MethodPost, Endpoint: "/cachedContents", Body: &createReq, }, &created) + now = time.Now() if err != nil || strings.TrimSpace(created.Name) == "" { if err == nil { err = fmt.Errorf("cached-content creation returned an empty name") diff --git a/internal/responsecache/exact_cache_test.go b/internal/responsecache/exact_cache_test.go index b8057908..6b96ba63 100644 --- a/internal/responsecache/exact_cache_test.go +++ b/internal/responsecache/exact_cache_test.go @@ -35,6 +35,9 @@ type blockingMissExchange struct { contextCalls atomic.Int32 } +// Context deliberately signals on StoreAfter's second Context call: the first +// reads the workflow and the second occurs after a follower joins the wait. +// Update this helper if StoreAfter's pre-wait call count changes. func (e *blockingMissExchange) Context() context.Context { if e.joined != nil && e.contextCalls.Add(1) == 2 { close(e.joined) From e52309054abba1eb2ae9a8fa70764c06c08c42f8 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 4 Aug 2026 21:30:07 +0200 Subject: [PATCH 5/6] fix(cache): harden coalescing and cache keys --- internal/providers/cache_planner.go | 28 +++-- internal/providers/cache_planner_test.go | 13 +++ internal/providers/gemini/gemini_test.go | 11 +- internal/responsecache/exact_cache_test.go | 122 +++++++++++++++++++-- internal/responsecache/simple.go | 31 ++++-- internal/responsecache/stream_cache.go | 57 ++++++++++ 6 files changed, 231 insertions(+), 31 deletions(-) diff --git a/internal/providers/cache_planner.go b/internal/providers/cache_planner.go index 9ee3e58f..4f2370fe 100644 --- a/internal/providers/cache_planner.go +++ b/internal/providers/cache_planner.go @@ -105,7 +105,12 @@ func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, sele func (p *cachePlanner) planResponses(req *core.ResponsesRequest, providerType string, selector core.ModelSelector) *core.ResponsesRequest { profile := promptCacheProfileFor(providerType) - if p == nil || !p.enabled || req == nil || profile.mode == promptCacheUnsupported { + if p == nil || !p.enabled || req == nil { + return req + } + // Only OpenAI and Anthropic Responses requests have cache directives that + // this planner can emit. Reject other modes before scanning or cloning. + if profile.mode != promptCacheOpenAI && profile.mode != promptCacheAnthropic { return req } items, ok := req.Input.([]core.ResponsesInputElement) @@ -171,8 +176,17 @@ func cloneResponsesRequest(req *core.ResponsesRequest) (*core.ResponsesRequest, return &clone, true } +var cacheDirectiveKeys = []string{ + "cache_control", + "cached_content", + "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_breakpoint", + core.GatewayCachePointField, +} + func hasCacheDirective(fields core.UnknownJSONFields) bool { - for _, key := range []string{"cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options", "prompt_cache_breakpoint", core.GatewayCachePointField} { + for _, key := range cacheDirectiveKeys { if len(fields.Lookup(key)) > 0 { return true } @@ -251,12 +265,7 @@ func anyHasCacheDirective(value any) bool { } func isCacheDirectiveKey(key string) bool { - switch key { - case "cache_control", "cached_content", "prompt_cache_key", "prompt_cache_options", "prompt_cache_breakpoint", core.GatewayCachePointField: - return true - default: - return false - } + return slices.Contains(cacheDirectiveKeys, key) } func estimateSimpleChatPrefixTokens(req *core.ChatRequest) (int, bool) { @@ -451,6 +460,9 @@ func providerCacheMinimum(profile promptCacheProfile, model string) int { return 1536 } if strings.Contains(model, "claude") { + // AWS documents a Bedrock-specific 4096-token checkpoint minimum + // for these models, including Sonnet 4.5; the direct Anthropic API + // minimum can differ for the same model family. if strings.Contains(model, "haiku-4-5") || strings.Contains(model, "sonnet-4-5") || strings.Contains(model, "opus-4-5") || strings.Contains(model, "opus-4-6") { return 4096 diff --git a/internal/providers/cache_planner_test.go b/internal/providers/cache_planner_test.go index a51d44dd..df039b55 100644 --- a/internal/providers/cache_planner_test.go +++ b/internal/providers/cache_planner_test.go @@ -217,6 +217,19 @@ func TestCachePlannerProviderCapabilityBoundaries(t *testing.T) { } } +func TestCachePlannerSkipsUnsupportedResponsesModesBeforeCloning(t *testing.T) { + req := &core.ResponsesRequest{Input: []core.ResponsesInputElement{ + {Role: "user", Content: strings.Repeat("x", 20000)}, + {Role: "user", Content: "turn"}, + }} + planner := &cachePlanner{enabled: true} + for _, provider := range []string{"bedrock", "gemini", "openrouter", "unknown"} { + if got := planner.planResponses(req, provider, core.ModelSelector{Model: "model"}); got != req { + t.Fatalf("provider %q unexpectedly received a Responses plan", provider) + } + } +} + func TestCloneChatRequestPreservesInternalCachePlan(t *testing.T) { req := &core.ChatRequest{PromptCachePlan: &core.PromptCachePlan{Key: "stable"}} clone, ok := cloneChatRequest(req) diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 436a41da..080aa32b 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -220,7 +220,13 @@ func TestPrepareCachedContentCoalescesConcurrentCreation(t *testing.T) { } func TestPrepareCachedContentRequiresStableCredentialAndAIStudio(t *testing.T) { - p := NewWithHTTPClient("key", http.DefaultClient, llmclient.Hooks{}) + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + requests.Add(1) + })) + defer server.Close() + p := NewWithHTTPClient("key", server.Client(), llmclient.Hooks{}) + p.SetBaseURL(server.URL) p.keys = providers.NewKeyring("one", "two") req := &core.ChatRequest{PromptCachePlan: &core.PromptCachePlan{Key: "prefix"}} body := &geminiGenerateContentRequest{SystemInstruction: &geminiContent{}, Contents: []geminiContent{{Role: "user"}}} @@ -238,6 +244,9 @@ func TestPrepareCachedContentRequiresStableCredentialAndAIStudio(t *testing.T) { if body.CachedContent != "" { t.Fatal("Vertex must not use the AI Studio cachedContents endpoint") } + if got := requests.Load(); got != 0 { + t.Fatalf("Vertex cache preparation made %d cachedContents requests, want none", got) + } } func TestGeminiCacheObjectMapIsBoundedAndSweepsExpiredEntries(t *testing.T) { diff --git a/internal/responsecache/exact_cache_test.go b/internal/responsecache/exact_cache_test.go index 6b96ba63..b60600b7 100644 --- a/internal/responsecache/exact_cache_test.go +++ b/internal/responsecache/exact_cache_test.go @@ -32,6 +32,7 @@ type blockingMissExchange struct { started chan<- struct{} release <-chan struct{} joined chan struct{} + nonCacheable bool contextCalls atomic.Int32 } @@ -59,6 +60,9 @@ func (e *blockingMissExchange) Capture(_ string, next func() error) ([]byte, boo if err := next(); err != nil { return nil, false, err } + if e.nonCacheable { + return nil, false, nil + } return []byte(`{"ok":true}`), true, nil } @@ -303,27 +307,103 @@ func TestStoreAfter_LeaderErrorIsNotFannedOut(t *testing.T) { }() <-started - joined := make(chan struct{}) + const followers = 8 + joined := make([]chan struct{}, followers) + errs := make([]error, followers) var followerCalls atomic.Int32 + var wg sync.WaitGroup + for i := range followers { + joined[i] = make(chan struct{}) + wg.Go(func() { + errs[i] = m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), joined: joined[i], + }, body, func() error { + followerCalls.Add(1) + return nil + }) + }) + } + for _, waiter := range joined { + <-waiter + } + close(release) + if err := <-leaderDone; !errors.Is(err, leaderErr) { + t.Fatalf("leader error = %v, want %v", err, leaderErr) + } + wg.Wait() + for i, err := range errs { + if err != nil { + t.Fatalf("follower %d inherited leader error: %v", i, err) + } + } + if got := followerCalls.Load(); got != followers { + t.Fatalf("follower provider calls = %d, want %d independent retries", got, followers) + } +} + +func TestStoreAfter_CacheableFollowerStoresAfterNonCacheableLeader(t *testing.T) { + store := cache.NewMapStore() + m := newSimpleCacheMiddleware(store, time.Hour, nil) + defer m.close() + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) + started := make(chan struct{}, 1) + release := make(chan struct{}) + leaderDone := make(chan error, 1) + go func() { + leaderDone <- m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), started: started, release: release, nonCacheable: true, + }, body, func() error { return nil }) + }() + <-started + + joined := make(chan struct{}) followerDone := make(chan error, 1) go func() { followerDone <- m.StoreAfter(&blockingMissExchange{ ctx: context.Background(), joined: joined, - }, body, func() error { - followerCalls.Add(1) - return nil - }) + }, body, func() error { return nil }) }() <-joined close(release) - if err := <-leaderDone; !errors.Is(err, leaderErr) { - t.Fatalf("leader error = %v, want %v", err, leaderErr) + if err := <-leaderDone; err != nil { + t.Fatalf("leader error: %v", err) } if err := <-followerDone; err != nil { - t.Fatalf("follower inherited leader error: %v", err) + t.Fatalf("follower error: %v", err) } - if got := followerCalls.Load(); got != 1 { - t.Fatalf("follower provider calls = %d, want independent retry", got) + m.wg.Wait() + key := hashRequest("/v1/chat/completions", body, nil) + if cached, err := store.Get(context.Background(), key); err != nil || len(cached) == 0 { + t.Fatalf("cached follower response = %q, err=%v", cached, err) + } +} + +func TestStoreAfter_LeaderPanicReleasesMiss(t *testing.T) { + m := newSimpleCacheMiddleware(cache.NewMapStore(), time.Hour, nil) + defer m.close() + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"same"}]}`) + var recovered any + func() { + defer func() { recovered = recover() }() + _ = m.StoreAfter(&blockingMissExchange{ctx: context.Background()}, body, func() error { + panic("provider panic") + }) + }() + if recovered == nil { + t.Fatal("leader panic did not propagate") + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + var calls atomic.Int32 + if err := m.StoreAfter(&blockingMissExchange{ctx: ctx}, body, func() error { + calls.Add(1) + return nil + }); err != nil { + t.Fatalf("request after leader panic: %v", err) + } + if got := calls.Load(); got != 1 { + t.Fatalf("provider calls after leader panic = %d, want 1", got) } } @@ -341,6 +421,8 @@ func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { {name: "array order matters", first: `{"input":[1,2],"model":"gpt-4"}`, second: `{"input":[2,1],"model":"gpt-4"}`, equal: false}, {name: "malformed falls back exactly", first: `{"input":`, second: ` {"input":`, equal: false}, {name: "multiple values fall back exactly", first: `{"a":1} {"b":2}`, second: `{"a":1} {"b":2}`, equal: false}, + {name: "duplicate names fall back exactly", first: `{"model":"a","model":"b"}`, second: `{"model":"b"}`, equal: false}, + {name: "nested duplicate names fall back exactly", first: `{"input":{"a":1,"a":2}}`, second: `{"input":{"a":2}}`, equal: false}, } { t.Run(tt.name, func(t *testing.T) { first := hashRequest("/v1/embeddings", []byte(tt.first), plan) @@ -352,6 +434,26 @@ func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { } } +func TestHashRequest_DuplicateNamesDoNotCollideAfterTypedDecoding(t *testing.T) { + plan := resolvedWorkflow("openai", "gpt-4") + for _, tt := range []struct { + path string + duplicate string + collapsed string + }{ + {path: "/v1/chat/completions", duplicate: `{"model":"a","model":"b","messages":[]}`, collapsed: `{"model":"b","messages":[]}`}, + {path: "/v1/responses", duplicate: `{"model":"a","model":"b","input":[]}`, collapsed: `{"model":"b","input":[]}`}, + } { + t.Run(tt.path, func(t *testing.T) { + first := hashRequest(tt.path, []byte(tt.duplicate), plan) + second := hashRequest(tt.path, []byte(tt.collapsed), plan) + if first == second { + t.Fatal("duplicate-member request collided with its collapsed form") + } + }) + } +} + func TestHashRequest_ResolvedModelChangesKey(t *testing.T) { body := []byte(`{"model":"anthropic/claude-opus-4-6","messages":[{"role":"user","content":"hi"}]}`) diff --git a/internal/responsecache/simple.go b/internal/responsecache/simple.go index 92efc89b..79320e3f 100644 --- a/internal/responsecache/simple.go +++ b/internal/responsecache/simple.go @@ -111,13 +111,10 @@ func (m *simpleCacheMiddleware) StoreAfter(ex exchange, body []byte, next func() call, leader := m.joinMiss(key) if leader { - data, ok, err := ex.Capture("response cache: failed to capture cacheable response body", next) - if err == nil && ok { - m.enqueueWrite(cacheWriteJob{key: key, data: data}) - } else { - data = nil - } - m.finishMiss(key, call, data, err) + var data []byte + var err error + defer func() { m.finishMiss(key, call, data, err) }() + data, err = m.captureAndStore(ex, key, next) return err } @@ -129,11 +126,9 @@ func (m *simpleCacheMiddleware) StoreAfter(ex exchange, body []byte, next func() // The leader produced a non-cacheable result (failure status, failover, or // malformed body). Waiting followers must execute independently rather than // replaying something the normal cache would refuse to store. - if call.err != nil { - return m.StoreAfter(ex, body, next) - } - if len(call.data) == 0 { - return next() + if call.err != nil || len(call.data) == 0 { + _, err := m.captureAndStore(ex, key, next) + return err } if err := ex.ReplayHit(body, call.data, CacheTypeExact); err != nil { return next() @@ -145,6 +140,18 @@ func (m *simpleCacheMiddleware) StoreAfter(ex exchange, body []byte, next func() return nil } +// captureAndStore executes one cache miss without joining the coalescing +// group. Followers use it after a leader produces no replayable response so +// they proceed independently while retaining the normal cache-write behavior. +func (m *simpleCacheMiddleware) captureAndStore(ex exchange, key string, next func() error) ([]byte, error) { + data, ok, err := ex.Capture("response cache: failed to capture cacheable response body", next) + if err != nil || !ok { + return nil, err + } + m.enqueueWrite(cacheWriteJob{key: key, data: data}) + return data, nil +} + func (m *simpleCacheMiddleware) joinMiss(key string) (*exactMissCall, bool) { m.missMu.Lock() defer m.missMu.Unlock() diff --git a/internal/responsecache/stream_cache.go b/internal/responsecache/stream_cache.go index 68ba3454..4d2dc772 100644 --- a/internal/responsecache/stream_cache.go +++ b/internal/responsecache/stream_cache.go @@ -2,6 +2,7 @@ package responsecache import ( "bytes" + stdjson "encoding/json" "io" "net/http" "strings" @@ -22,6 +23,12 @@ var ( ) func cacheKeyRequestBody(path string, body []byte) []byte { + // Typed decoding and generic unmarshaling both collapse duplicate object + // names. Keep such requests byte-exact so they cannot collide with the + // single-member form a provider may interpret differently. + if hasDuplicateJSONMemberNames(body) { + return body + } switch path { case "/v1/chat/completions": req, err := core.DecodeChatRequest(body, nil) @@ -58,6 +65,56 @@ func cacheKeyRequestBody(path string, body []byte) []byte { } } +func hasDuplicateJSONMemberNames(body []byte) bool { + decoder := stdjson.NewDecoder(bytes.NewReader(body)) + var scanValue func() (bool, error) + scanValue = func() (bool, error) { + token, err := decoder.Token() + if err != nil { + return false, err + } + delim, ok := token.(stdjson.Delim) + if !ok { + return false, nil + } + switch delim { + case '{': + seen := make(map[string]struct{}) + for decoder.More() { + keyToken, err := decoder.Token() + if err != nil { + return false, err + } + key, ok := keyToken.(string) + if !ok { + return false, io.ErrUnexpectedEOF + } + if _, duplicate := seen[key]; duplicate { + return true, nil + } + seen[key] = struct{}{} + if duplicate, err := scanValue(); err != nil || duplicate { + return duplicate, err + } + } + _, err = decoder.Token() + return false, err + case '[': + for decoder.More() { + if duplicate, err := scanValue(); err != nil || duplicate { + return duplicate, err + } + } + _, err = decoder.Token() + return false, err + default: + return false, nil + } + } + duplicate, err := scanValue() + return err == nil && duplicate +} + // canonicalJSONForCache makes semantically identical JSON bodies share an // exact-cache key despite insignificant whitespace or object-key ordering. It // preserves number spellings through json.Number and falls back byte-for-byte From dc9ce66900ff81b1ff847c312f852dff783759a3 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 4 Aug 2026 21:45:13 +0200 Subject: [PATCH 6/6] fix(cache): preserve duplicate keys after large numbers --- internal/responsecache/exact_cache_test.go | 1 + internal/responsecache/stream_cache.go | 1 + 2 files changed, 2 insertions(+) diff --git a/internal/responsecache/exact_cache_test.go b/internal/responsecache/exact_cache_test.go index b60600b7..3826328b 100644 --- a/internal/responsecache/exact_cache_test.go +++ b/internal/responsecache/exact_cache_test.go @@ -423,6 +423,7 @@ func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { {name: "multiple values fall back exactly", first: `{"a":1} {"b":2}`, second: `{"a":1} {"b":2}`, equal: false}, {name: "duplicate names fall back exactly", first: `{"model":"a","model":"b"}`, second: `{"model":"b"}`, equal: false}, {name: "nested duplicate names fall back exactly", first: `{"input":{"a":1,"a":2}}`, second: `{"input":{"a":2}}`, equal: false}, + {name: "oversized number before duplicate names", first: `{"n":1e1000000,"model":"a","model":"b"}`, second: `{"n":1e1000000,"model":"b"}`, equal: false}, } { t.Run(tt.name, func(t *testing.T) { first := hashRequest("/v1/embeddings", []byte(tt.first), plan) diff --git a/internal/responsecache/stream_cache.go b/internal/responsecache/stream_cache.go index 4d2dc772..219ad232 100644 --- a/internal/responsecache/stream_cache.go +++ b/internal/responsecache/stream_cache.go @@ -67,6 +67,7 @@ func cacheKeyRequestBody(path string, body []byte) []byte { func hasDuplicateJSONMemberNames(body []byte) bool { decoder := stdjson.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() var scanValue func() (bool, error) scanValue = func() (bool, error) { token, err := decoder.Token()