diff --git a/.env.template b/.env.template index d10e1821..49530327 100644 --- a/.env.template +++ b/.env.template @@ -117,6 +117,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; empty +# or invalid values keep the default. Client cache directives still pass through. +# 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 e44d9cb7..2b66e14a 100644 --- a/docs/features/cache.mdx +++ b/docs/features/cache.mdx @@ -90,6 +90,40 @@ 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 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 +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 +``` + +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/pro.mdx b/docs/pro.mdx index a6a9fad0..991b3fde 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 @@ -104,12 +110,20 @@ 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 | | `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/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 2709e01c..b4ec8e45 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -41,6 +41,18 @@ 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"` +} + +// 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 } func (r *ChatRequest) semanticSelector() (string, 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 9fd6e3a1..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 { @@ -153,6 +220,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 +234,21 @@ 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 } + // 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) default: return nil, nil, core.NewInvalidRequestError("unsupported message role: "+msg.Role, nil) @@ -181,6 +259,11 @@ func convertMessages(messages []core.Message) ([]brtypes.SystemContentBlock, []b return system, out, nil } +func isGatewayCachePoint(fields core.UnknownJSONFields) bool { + var enabled bool + return json.Unmarshal(fields.Lookup(core.GatewayCachePointField), &enabled) == nil && enabled +} + 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/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 new file mode 100644 index 00000000..4f2370fe --- /dev/null +++ b/internal/providers/cache_planner.go @@ -0,0 +1,511 @@ +package providers + +import ( + "crypto/sha256" + "encoding/hex" + "log/slog" + "os" + "slices" + "strconv" + "strings" + + "github.com/goccy/go-json" + + "github.com/enterpilot/gomodel/internal/core" +) + +const ( + providerPromptCachePlannerEnabledEnv = "PROVIDER_PROMPT_CACHE_PLANNER_ENABLED" +) + +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{enabled: true} + } + 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{enabled: true} + } + return &cachePlanner{enabled: enabled} +} + +func (p *cachePlanner) planChat(req *core.ChatRequest, providerType string, selector core.ModelSelector) *core.ChatRequest { + 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, prefixMessages}) + if err != nil || estimatedTokens(prefixBody) < minimum { + return req + } + + planned, ok := cloneChatRequest(req) + if !ok { + return req + } + key := cacheAffinityKey(providerType, selector, req.User, prefixBody) + switch profile.mode { + case promptCacheOpenAI: + 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 promptCacheAnthropic: + planned.ExtraFields = mergeCacheExtras(planned.ExtraFields, map[string]json.RawMessage{ + "cache_control": json.RawMessage(`{"type":"ephemeral"}`), + }) + 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 { + profile := promptCacheProfileFor(providerType) + 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) + if !ok || len(items) < 2 || hasResponsesCacheDirective(req, items[:len(items)-1]) { + 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(profile, selector.Model) { + return req + } + planned, ok := cloneResponsesRequest(req) + if !ok { + return req + } + key := cacheAffinityKey(providerType, selector, req.User, prefixBody) + switch profile.mode { + case promptCacheOpenAI: + 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 promptCacheAnthropic: + 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 + } + if req.PromptCachePlan != nil { + plan := *req.PromptCachePlan + clone.PromptCachePlan = &plan + } + 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 +} + +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 cacheDirectiveKeys { + if len(fields.Lookup(key)) > 0 { + return true + } + } + 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 { + return slices.Contains(cacheDirectiveKeys, key) +} + +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 + } + 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 == "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"}`), + }) + 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 == "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"} + } + items[i].Content = content + req.Input = items + 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 +} + +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(profile promptCacheProfile, model string) int { + model = strings.NewReplacer(".", "-", "_", "-").Replace(strings.ToLower(model)) + switch profile.mode { + case promptCacheOpenAI: + return 1024 + case promptCacheAnthropic: + 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") { + return 2048 + } + return 1024 + case promptCacheGemini: + return 4096 + case promptCacheBedrock: + if strings.Contains(model, "nova") { + 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 + } + 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 "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: + return promptCacheProfile{} + } +} + +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..df039b55 --- /dev/null +++ b/internal/providers/cache_planner_test.go @@ -0,0 +1,273 @@ +package providers + +import ( + "bytes" + "os" + "strings" + "testing" + + "github.com/goccy/go-json" + + "github.com/enterpilot/gomodel/internal/core" +) + +func TestCachePlannerAppliesProviderSpecificChatPlanWithoutMutatingCaller(t *testing.T) { + planner := &cachePlanner{enabled: true} + 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: core.GatewayCachePointField}, + {provider: "bedrock-mantle", model: "gpt-5.6", field: "prompt_cache_key"}, + {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 := planner.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 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) + } + } + if !req.ExtraFields.IsEmpty() || !req.Messages[0].ExtraFields.IsEmpty() { + t.Fatal("planner mutated caller-owned request") + } + }) + } +} + +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 + value string + set bool + enabled bool + }{ + {name: "default on", 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) { + 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) + } + }) + } +} + +func TestCachePlannerHonorsMinimumAndClientDirective(t *testing.T) { + 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") + } + + 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") + } +} + +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 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) + 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") + } +} + +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 b2a33664..a6764518 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -3,16 +3,21 @@ package gemini import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "io" + "log/slog" "net/http" "net/url" "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" @@ -61,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 @@ -74,6 +84,15 @@ 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 + retryAfter time.Time } // New creates a new Gemini provider. @@ -462,6 +481,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 +529,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 +542,148 @@ 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 != "" || p.backend != geminiBackendAIStudio || len(body.Contents) == 0 { + return + } + if req.PromptCachePlan == nil || strings.TrimSpace(req.PromptCachePlan.Key) == "" { + return + } + // 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, _, _ := 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, + Contents: append([]geminiContent(nil), body.Contents[:len(body.Contents)-1]...), + Tools: append([]geminiTool(nil), body.Tools...), + 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") + } + 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 = now.Add(geminiCacheTTL) + } + if !now.Add(geminiCacheFreshness).Before(expiresAt) { + p.storeCachedContentObject(key, geminiCacheObject{retryAfter: now.Add(geminiCacheFailureBackoff)}, now) + return "", nil + } + p.storeCachedContentObject(key, geminiCacheObject{name: created.Name, expiresAt: expiresAt}, now) + return created.Name, nil + }) + 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, now time.Time) (string, bool) { + p.cacheMu.Lock() + defer p.cacheMu.Unlock() + entry, ok := p.cacheObjects[key] + if !ok { + return "", false + } + if entry.name == "" { + if now.Before(entry.retryAfter) { + return "", true + } + delete(p.cacheObjects, key) + return "", false + } + if now.Add(geminiCacheFreshness).After(entry.expiresAt) { + delete(p.cacheObjects, key) + 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 + } + } + p.cacheObjects[key] = entry +} + +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..080aa32b 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -1,14 +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" @@ -34,6 +39,233 @@ 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" { + wrongPath.Store(true) + http.Error(w, "wrong path", http.StatusNotFound) + return + } + 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"}}}, + }, + Tools: []geminiTool{{FunctionDeclarations: []geminiFunctionDeclaration{{Name: "lookup"}}}}, + } + } + 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()) + } + 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) + } + } +} + +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) { + 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"}}} + 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") + } + if got := requests.Load(); got != 0 { + t.Fatalf("Vertex cache preparation made %d cachedContents requests, want none", got) + } +} + +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 128901cc..b0ecf3cb 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,22 @@ 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 { + 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 + } + 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 +656,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 +665,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 +706,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 +715,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..3826328b 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,45 @@ type concurrentTrackingStore struct { releaseCh chan struct{} } +type blockingMissExchange struct { + ctx context.Context + started chan<- struct{} + release <-chan struct{} + joined chan struct{} + nonCacheable bool + 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) + } + 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 + } + if e.nonCacheable { + return nil, false, nil + } + return []byte(`{"ok":true}`), true, nil +} + func newConcurrentTrackingStore() *concurrentTrackingStore { return &concurrentTrackingStore{ enterCh: make(chan struct{}, 1024), @@ -86,6 +126,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 +151,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,6 +217,244 @@ func TestHandleRequest_DifferentBodyDifferentKey(t *testing.T) { } } +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{}, 1) + release := make(chan struct{}) + const requests = 12 + errs := make([]error, requests) + var wg sync.WaitGroup + 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() { + errs[i] = m.StoreAfter(&blockingMissExchange{ + ctx: context.Background(), joined: joined[i-1], + }, body, func() error { + calls.Add(1) + return nil + }) + }) + } + 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) + } + for i := range requests { + if errs[i] != nil { + t.Fatalf("request %d: %v", i, errs[i]) + } + } +} + +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 + + 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 { return nil }) + }() + <-joined + close(release) + if err := <-leaderDone; err != nil { + t.Fatalf("leader error: %v", err) + } + if err := <-followerDone; err != nil { + t.Fatalf("follower error: %v", err) + } + 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) + } +} + +func TestHashRequest_CanonicalizesJSONFormattingAndKeyOrder(t *testing.T) { + plan := resolvedWorkflow("openai", "gpt-4") + 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}, + {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) + 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) + } + }) + } +} + +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 e2cd176f..79320e3f 100644 --- a/internal/responsecache/simple.go +++ b/internal/responsecache/simple.go @@ -47,6 +47,14 @@ type simpleCacheMiddleware struct { workers sync.WaitGroup mu sync.RWMutex closed bool + 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 { @@ -101,17 +109,73 @@ 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) - if err != nil { + call, leader := m.joinMiss(key) + if leader { + var data []byte + var err error + defer func() { m.finishMiss(key, call, data, err) }() + data, err = m.captureAndStore(ex, key, next) return err } - if !ok { - 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 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() + } + ex.MarkHit(CacheTypeExact) + if m.hitRecorder != nil { + m.hitRecorder(ex, call.data, CacheTypeExact) } - m.enqueueWrite(cacheWriteJob{key: key, data: data}) 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() + 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() diff --git a/internal/responsecache/stream_cache.go b/internal/responsecache/stream_cache.go index 6353fd1e..219ad232 100644 --- a/internal/responsecache/stream_cache.go +++ b/internal/responsecache/stream_cache.go @@ -2,6 +2,8 @@ package responsecache import ( "bytes" + stdjson "encoding/json" + "io" "net/http" "strings" @@ -21,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) @@ -53,8 +61,80 @@ func cacheKeyRequestBody(path string, body []byte) []byte { } return normalized default: + return canonicalJSONForCache(body) + } +} + +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() + 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 +// 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 {