From 179bf2f93c6433fddaa4d404ef1c71607399c730 Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:20:58 -0500 Subject: [PATCH 1/2] fix(go/plugins/googlegenai): do not resend the cached prefix inline handleCache already stores messages[0..endIndex] on the CachedContent resource. toGeminiContents still sent those same messages in the generate request, so the prefix was billed twice. Skip the cached span when building contents, matching the JS plugin. Fixes #6137 --- go/plugins/googlegenai/cache.go | 27 +++++++------- go/plugins/googlegenai/cache_test.go | 51 +++++++++++++++++++++++++++ go/plugins/googlegenai/gemini.go | 16 ++++++--- go/plugins/googlegenai/gemini_test.go | 2 +- 4 files changed, 79 insertions(+), 17 deletions(-) diff --git a/go/plugins/googlegenai/cache.go b/go/plugins/googlegenai/cache.go index ded7e6755f..d77d03383e 100644 --- a/go/plugins/googlegenai/cache.go +++ b/go/plugins/googlegenai/cache.go @@ -46,33 +46,33 @@ func handleCache( client *genai.Client, request *ai.ModelRequest, model string, -) (*genai.CachedContent, error) { +) (*genai.CachedContent, int, error) { cs, err := findCacheMarker(request) if err != nil { - return nil, err + return nil, -1, err } if cs == nil { - return nil, nil + return nil, -1, nil } // no cache mark found if cs.endIndex == -1 { - return nil, err + return nil, -1, err } // index out of bounds if cs.endIndex < 0 || cs.endIndex >= len(request.Messages) { - return nil, status.Errorf(status.ErrInvalidArgument, "end of cached contents, index %d is invalid", cs.endIndex) + return nil, -1, status.Errorf(status.ErrInvalidArgument, "end of cached contents, index %d is invalid", cs.endIndex) } // since context caching is only available for specific model versions, we // must make sure the configuration has the right version err = validateContextCacheRequest(request, model) if err != nil { - return nil, err + return nil, -1, err } messages, err := messagesToCache(request.Messages, cs.endIndex) if err != nil { - return nil, err + return nil, -1, err } hash := calculateCacheHash(messages) @@ -81,14 +81,14 @@ func handleCache( cache, err = lookupCache(ctx, client, cs.name) if err != nil { // TODO: if cache expired or not found, create a fresh one - return nil, fmt.Errorf("cache lookup error: %w", wrapAPIError(err)) + return nil, -1, fmt.Errorf("cache lookup error: %w", wrapAPIError(err)) } // make sure the cache contents matches the request messages hash if cache.DisplayName != hash { - return nil, status.Errorf(status.ErrInvalidArgument, "invalid cache name: hash mismatch between cached content and request messages") + return nil, -1, status.Errorf(status.ErrInvalidArgument, "invalid cache name: hash mismatch between cached content and request messages") } - return cache, nil + return cache, cs.endIndex, nil } if cs.ttl > 0 { @@ -98,11 +98,14 @@ func handleCache( Contents: messages, }) if err != nil { - return nil, fmt.Errorf("cache creation error: %w", wrapAPIError(err)) + return nil, -1, fmt.Errorf("cache creation error: %w", wrapAPIError(err)) } } - return cache, nil + if cache == nil { + return nil, -1, nil + } + return cache, cs.endIndex, nil } // messagesToCache collects all the messages that should be cached diff --git a/go/plugins/googlegenai/cache_test.go b/go/plugins/googlegenai/cache_test.go index 7be0cef2c9..67e71b275b 100644 --- a/go/plugins/googlegenai/cache_test.go +++ b/go/plugins/googlegenai/cache_test.go @@ -198,6 +198,57 @@ func TestExtractCacheConfig_InvalidCacheType(t *testing.T) { } } +func TestToGeminiContents_DropsCachedPrefix(t *testing.T) { + prefix := "this is a large stable prefix that should only live in the cache" + followUp := "what did I just say?" + req := &ai.ModelRequest{ + Messages: []*ai.Message{ + { + Role: ai.RoleUser, + Content: []*ai.Part{{Text: prefix}}, + Metadata: map[string]any{ + "cache": map[string]any{"ttlSeconds": 3600}, + }, + }, + { + Role: ai.RoleUser, + Content: []*ai.Part{{Text: followUp}}, + }, + }, + } + + cs, err := findCacheMarker(req) + if err != nil { + t.Fatalf("findCacheMarker: %v", err) + } + if cs == nil || cs.endIndex != 0 { + t.Fatalf("cache marker = %#v, want endIndex=0", cs) + } + + // Reproduce #6137: sending the whole request inline includes the prefix. + inline, err := toGeminiContents(req, -1) + if err != nil { + t.Fatalf("toGeminiContents(no skip): %v", err) + } + if len(inline) != 2 { + t.Fatalf("inline contents = %d, want 2 (bug: prefix sent with the follow-up)", len(inline)) + } + if got := inline[0].Parts[0].Text; got != prefix { + t.Fatalf("inline[0] = %q, want cached prefix", got) + } + + got, err := toGeminiContents(req, cs.endIndex) + if err != nil { + t.Fatalf("toGeminiContents(skip cached): %v", err) + } + if len(got) != 1 { + t.Fatalf("contents after cache boundary = %d, want 1", len(got)) + } + if got[0].Parts[0].Text != followUp { + t.Errorf("remaining content = %q, want %q", got[0].Parts[0].Text, followUp) + } +} + func TestFindCacheMarker_NumericTTLForms(t *testing.T) { // ttlSeconds is an int when set from Go code but arrives as float64 or // json.Number after a JSON round-trip (dev UI, reflection server). All diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 4cd231257c..f41b35af6d 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -159,7 +159,7 @@ func generate( } model = resolveVertexModelName(client, model) - cache, err := handleCache(ctx, client, input, model) + cache, cachedThrough, err := handleCache(ctx, client, input, model) if err != nil { return nil, err } @@ -169,7 +169,7 @@ func generate( return nil, err } - contents, err := toGeminiContents(input) + contents, err := toGeminiContents(input, cachedThrough) if err != nil { return nil, err } @@ -324,9 +324,17 @@ func mergeCandidateMetadata(dst, src *genai.Candidate) { // toGeminiContents converts the non-system messages of an [*ai.ModelRequest] // to a slice of [*genai.Content]. System messages are handled separately via // the request's system instruction. -func toGeminiContents(input *ai.ModelRequest) ([]*genai.Content, error) { +// +// cachedThrough is the inclusive index of the last message stored in a +// CachedContent resource, or -1 if no cache is in use. Messages at or before +// that index must not be sent inline — the JS plugin slices the cached span +// off the request, and sending it here would bill the prefix twice (#6137). +func toGeminiContents(input *ai.ModelRequest, cachedThrough int) ([]*genai.Content, error) { var contents []*genai.Content - for _, m := range input.Messages { + for i, m := range input.Messages { + if cachedThrough >= 0 && i <= cachedThrough { + continue + } // system parts are handled separately if m.Role == ai.RoleSystem { continue diff --git a/go/plugins/googlegenai/gemini_test.go b/go/plugins/googlegenai/gemini_test.go index 2e63d762ba..4c3dd0a334 100644 --- a/go/plugins/googlegenai/gemini_test.go +++ b/go/plugins/googlegenai/gemini_test.go @@ -1438,7 +1438,7 @@ func TestToGeminiContents(t *testing.T) { }, } - contents, err := toGeminiContents(input) + contents, err := toGeminiContents(input, -1) if err != nil { t.Fatalf("toGeminiContents failed: %v", err) } From c19ca65ed13ebeac8b2554e9a1202f027158f3bf Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:20:58 -0500 Subject: [PATCH 2/2] fix(go/plugins/googlegenai): cache contents in chronological order messagesToCache walked the prefix backwards, so after we stop sending the prefix inline the model would only see a reversed cache. --- go/plugins/googlegenai/cache.go | 15 ++++++++------- go/plugins/googlegenai/cache_test.go | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/go/plugins/googlegenai/cache.go b/go/plugins/googlegenai/cache.go index d77d03383e..95696cf2f4 100644 --- a/go/plugins/googlegenai/cache.go +++ b/go/plugins/googlegenai/cache.go @@ -39,8 +39,9 @@ var invalidArgMessages = struct { systemPrompt: "system prompts are not supported with context caching", } -// handleCache checks if caching should be used, attempts to find or create the cache, -// and returns the cached content if applicable. +// handleCache checks if caching should be used, attempts to find or create the +// cache, and returns the cached content plus the inclusive index of the last +// cached message (-1 when no cache is in use). func handleCache( ctx context.Context, client *genai.Client, @@ -111,18 +112,18 @@ func handleCache( // messagesToCache collects all the messages that should be cached func messagesToCache(m []*ai.Message, cacheEndIdx int) ([]*genai.Content, error) { var messagesToCache []*genai.Content - for i := cacheEndIdx; i >= 0; i-- { - m := m[i] - if m.Role == ai.RoleSystem { + for i := 0; i <= cacheEndIdx; i++ { + msg := m[i] + if msg.Role == ai.RoleSystem { continue } - parts, err := toGeminiParts(m.Content) + parts, err := toGeminiParts(msg.Content) if err != nil { return nil, err } messagesToCache = append(messagesToCache, &genai.Content{ Parts: parts, - Role: string(m.Role), + Role: string(msg.Role), }) } return messagesToCache, nil diff --git a/go/plugins/googlegenai/cache_test.go b/go/plugins/googlegenai/cache_test.go index 67e71b275b..501144e489 100644 --- a/go/plugins/googlegenai/cache_test.go +++ b/go/plugins/googlegenai/cache_test.go @@ -249,6 +249,24 @@ func TestToGeminiContents_DropsCachedPrefix(t *testing.T) { } } +func TestMessagesToCache_PreservesChronologicalOrder(t *testing.T) { + req := []*ai.Message{ + {Role: ai.RoleUser, Content: []*ai.Part{{Text: "first"}}}, + {Role: ai.RoleModel, Content: []*ai.Part{{Text: "second"}}}, + {Role: ai.RoleUser, Content: []*ai.Part{{Text: "third, not cached"}}}, + } + got, err := messagesToCache(req, 1) + if err != nil { + t.Fatalf("messagesToCache: %v", err) + } + if len(got) != 2 { + t.Fatalf("len = %d, want 2", len(got)) + } + if got[0].Parts[0].Text != "first" || got[1].Parts[0].Text != "second" { + t.Errorf("order = %q, %q; want first, second", got[0].Parts[0].Text, got[1].Parts[0].Text) + } +} + func TestFindCacheMarker_NumericTTLForms(t *testing.T) { // ttlSeconds is an int when set from Go code but arrives as float64 or // json.Number after a JSON round-trip (dev UI, reflection server). All