diff --git a/go/plugins/googlegenai/cache.go b/go/plugins/googlegenai/cache.go index ded7e6755f..95696cf2f4 100644 --- a/go/plugins/googlegenai/cache.go +++ b/go/plugins/googlegenai/cache.go @@ -39,40 +39,41 @@ 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, 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 +82,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,28 +99,31 @@ 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 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 7be0cef2c9..501144e489 100644 --- a/go/plugins/googlegenai/cache_test.go +++ b/go/plugins/googlegenai/cache_test.go @@ -198,6 +198,75 @@ 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 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 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) }