Repository: RooVetGit/Roo-Code
Pull request #2925: Gemini caching improvements
Author: login: cte
Labels: name: Enhancement
description: New feature or request
name: size:L
description: This PR changes 100-499 lines, ignoring generated files.
name: lgtm
description: This PR has been approved by a maintainer
Created: 2025-04-24T20:12:57Z | Merged: 2025-04-24T21:20:10Z

## Context

- Make sure there's enough content to cache
- Don't block on cache writes
- Ignore cache write errors
- Use a `NodeCache` instead of a `Map` to expire keys after 5m
- Flip caching back on for supported Gemini models
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Improves Gemini caching by using `NodeCache`, handling cache errors, and re-enabling caching for supported models.
> 
>   - **Caching Improvements**:
>     - Replaced `Map` with `NodeCache` in `GeminiHandler` for expiring keys after 5 minutes.
>     - Added `isCacheBusy` flag to prevent blocking on cache writes in `gemini.ts`.
>     - Implemented error handling for cache write failures in `GeminiHandler`.
>   - **Behavior Changes**:
>     - Enabled prompt caching for supported Gemini models in `api.ts`.
>     - Added minimum token count check for caching in `createMessage()` in `gemini.ts`.
>   - **Miscellaneous**:
>     - Added `getMessagesLength()` function in `gemini-format.ts` to calculate message length.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=RooVetGit%2FRoo-Code&utm_source=github&utm_medium=referral)<sup> for dfd7bb3926c7d081e5ede1e33019033cc1e7a9ac. You can [customize](https://app.ellipsis.dev/RooVetGit/settings/summaries) this summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->


Diff:
diff --git a/.changeset/bright-singers-drop.md b/.changeset/bright-singers-drop.md
new file mode 100644
index 00000000000..f76b8001b84
--- /dev/null
+++ b/.changeset/bright-singers-drop.md
@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+Enable prompt caching for Gemini (with some improvements)
diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts
index 43fae541379..ecc493dcd95 100644
--- a/src/api/providers/gemini.ts
+++ b/src/api/providers/gemini.ts
@@ -4,27 +4,42 @@ import {
 	type GenerateContentResponseUsageMetadata,
 	type GenerateContentParameters,
 	type Content,
+	CreateCachedContentConfig,
 } from "@google/genai"
+import NodeCache from "node-cache"
 
 import { SingleCompletionHandler } from "../"
 import type { ApiHandlerOptions, GeminiModelId, ModelInfo } from "../../shared/api"
 import { geminiDefaultModelId, geminiModels } from "../../shared/api"
-import { convertAnthropicContentToGemini, convertAnthropicMessageToGemini } from "../transform/gemini-format"
+import {
+	convertAnthropicContentToGemini,
+	convertAnthropicMessageToGemini,
+	getMessagesLength,
+} from "../transform/gemini-format"
 import type { ApiStream } from "../transform/stream"
 import { BaseProvider } from "./base-provider"
 
 const CACHE_TTL = 5
 
+const CONTEXT_CACHE_TOKEN_MINIMUM = 4096
+
+type CacheEntry = {
+	key: string
+	count: number
+}
+
 export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
+
 	private client: GoogleGenAI
-	private contentCaches: Map<string, { key: string; count: number }>
+	private contentCaches: NodeCache
+	private isCacheBusy = false
 
 	constructor(options: ApiHandlerOptions) {
 		super()
 		this.options = options
 		this.client = new GoogleGenAI({ apiKey: options.geminiApiKey ?? "not-provided" })
-		this.contentCaches = new Map()
+		this.contentCaches = new NodeCache({ stdTTL: 5 * 60, checkperiod: 5 * 60 })
 	}
 
 	async *createMessage(
@@ -35,36 +50,76 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
+
 		let uncachedContent: Content[] | undefined = undefined
 		let cachedContent: string | undefined = undefined
 		let cacheWriteTokens: number | undefined = undefined
 
+		// The minimum input token count for context caching is 4,096.
+		// For a basic approximation we assume 4 characters per token.
+		// We can use tiktoken eventually to get a more accurat token count.
 		// https://ai.google.dev/gemini-api/docs/caching?lang=node
-		// if (info.supportsPromptCache && cacheKey) {
-		// 	const cacheEntry = this.contentCaches.get(cacheKey)
-
-		// 	if (cacheEntry) {
-		// 		uncachedContent = contents.slice(cacheEntry.count, contents.length)
-		// 		cachedContent = cacheEntry.key
-		// 	}
+		// https://ai.google.dev/gemini-api/docs/tokens?lang=node
+		const isCacheAvailable =
+			info.supportsPromptCache &&
+			this.options.promptCachingEnabled &&
+			cacheKey &&
+			contentsLength > 4 * CONTEXT_CACHE_TOKEN_MINIMUM
+
+		if (isCacheAvailable) {
+			const cacheEntry = this.contentCaches.get<CacheEntry>(cacheKey)
+
+			if (cacheEntry) {
+				uncachedContent = contents.slice(cacheEntry.count, contents.length)
+				cachedContent = cacheEntry.key
+				console.log(
+					`[GeminiHandler] using ${cacheEntry.count} cached messages (${cacheEntry.key}) and ${uncachedContent.length} uncached messages`,
+				)
+			}
 
-		// 	const newCacheEntry = await this.client.caches.create({
-		// 		model,
-		// 		config: { contents, systemInstruction, ttl: `${CACHE_TTL * 60}s` },
-		// 	})
+			if (!this.isCacheBusy) {
+				this.isCacheBusy = true
+				const timestamp = Date.now()
+
+				this.client.caches
+					.create({
+						model,
+						config: {
+							contents,
+							systemInstruction,
+							ttl: `${CACHE_TTL * 60}s`,
+							httpOptions: { timeout: 120_000 },
+						},
+					})
+					.then((result) => {
+						const { name, usageMetadata } = result
+
+						if (name) {
+							this.contentCaches.set<CacheEntry>(cacheKey, { key: name, count: contents.length })
+							cacheWriteTokens = usageMetadata?.totalTokenCount ?? 0
+							console.log(
+								`[GeminiHandler] cached ${contents.length} messages (${cacheWriteTokens} tokens) in ${Date.now() - timestamp}ms`,
+							)
+						}
+					})
+					.catch((error) => {
+						console.error(`[GeminiHandler] caches.create error`, error)
+					})
+					.finally(() => {
+						this.isCacheBusy = false
+					})
+			}
+		}
 
-		// 	if (newCacheEntry.name) {
-		// 		this.contentCaches.set(cacheKey, { key: newCacheEntry.name, count: contents.length })
-		// 		cacheWriteTokens = newCacheEntry.usageMetadata?.totalTokenCount ?? 0
-		// 	}
-		// }
+		const isCacheUsed = !!cachedContent
 
 		const params: GenerateContentParameters = {
 			model,
 			contents: uncachedContent ?? contents,
 			config: {
 				cachedContent,
-				systemInstruction: cachedContent ? undefined : systemInstruction,
+				systemInstruction: isCacheUsed ? undefined : systemInstruction,
 				httpOptions: this.options.googleGeminiBaseUrl
 					? { baseUrl: this.options.googleGeminiBaseUrl }
 					: undefined,
@@ -94,13 +149,15 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			const cacheReadTokens = lastUsageMetadata.cachedContentTokenCount
 			const reasoningTokens = lastUsageMetadata.thoughtsTokenCount
 
-			// const totalCost = this.calculateCost({
-			// 	info,
-			// 	inputTokens,
-			// 	outputTokens,
-			// 	cacheWriteTokens,
-			// 	cacheReadTokens,
-			// })
+			const totalCost = isCacheUsed
+				? this.calculateCost({
+						info,
+						inputTokens,
+						outputTokens,
+						cacheWriteTokens,
+						cacheReadTokens,
+					})
+				: undefined
 
 			yield {
 				type: "usage",
@@ -109,7 +166,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 				cacheWriteTokens,
 				cacheReadTokens,
 				reasoningTokens,
-				// totalCost,
+				totalCost,
 			}
 		}
 	}
diff --git a/src/api/transform/gemini-format.ts b/src/api/transform/gemini-format.ts
index ee22cff32a4..be08d7ff7ba 100644
--- a/src/api/transform/gemini-format.ts
+++ b/src/api/transform/gemini-format.ts
@@ -76,3 +76,9 @@ export function convertAnthropicMessageToGemini(message: Anthropic.Messages.Mess
 		parts: convertAnthropicContentToGemini(message.content),
 	}
 }
+
+const getContentLength = ({ parts }: Content): number =>
+	parts?.reduce((length, { text }) => length + (text?.length ?? 0), 0) ?? 0
+
+export const getMessagesLength = (contents: Content[]): number =>
+	contents.reduce((length, content) => length + getContentLength(content), 0)
diff --git a/src/shared/api.ts b/src/shared/api.ts
index 25d4a668526..5dbed396a08 100644
--- a/src/shared/api.ts
+++ b/src/shared/api.ts
@@ -679,7 +679,7 @@ export const geminiModels = {
 		maxTokens: 65_535,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
 		isPromptCacheOptional: true,
 		inputPrice: 2.5, // This is the pricing for prompts above 200k tokens.
 		outputPrice: 15,
@@ -704,7 +704,7 @@ export const geminiModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
 		isPromptCacheOptional: true,
 		inputPrice: 0.1,
 		outputPrice: 0.4,
@@ -755,7 +755,7 @@ export const geminiModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
 		isPromptCacheOptional: true,
 		inputPrice: 0.15, // This is the pricing for prompts above 128k tokens.
 		outputPrice: 0.6,


Reviews:
user: login: mrubens
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-24T20:15:25Z
user: login: ellipsis-dev[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T20:15:41Z
user: login: cte
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-24T20:53:25Z
user: login: ellipsis-dev[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T21:02:33Z
user: login: ellipsis-dev[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T21:21:12Z

Review comments:
diff_hunk: @@ -35,36 +46,65 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
path: src/api/providers/gemini.ts
user: login: ellipsis-dev[bot]
body: The caching condition uses a raw character length (`systemInstruction.length + getMessagesLength(contents)`) to decide if caching is available. Consider using a proper tokenization method if Gemini expects token counts instead of character counts.

created_at: 2025-04-24T20:15:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,36 +46,65 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
+
 		let uncachedContent: Content[] | undefined = undefined
 		let cachedContent: string | undefined = undefined
 		let cacheWriteTokens: number | undefined = undefined
 
+		const isCacheAvailable =
+			info.supportsPromptCache && this.options.promptCachingEnabled && cacheKey && contentsLength > 16_384
path: src/api/providers/gemini.ts
user: login: ellipsis-dev[bot]
body: Magic number `16_384` is used in the cache availability condition. Consider extracting it to a named constant for clarity and easier adjustment.

created_at: 2025-04-24T20:15:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,36 +46,65 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
path: src/api/providers/gemini.ts
user: login: cte
body: Good call, but this is just an approximation for now.
created_at: 2025-04-24T20:53:25Z
author_association: COLLABORATOR
diff_hunk: @@ -35,36 +48,72 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
+
 		let uncachedContent: Content[] | undefined = undefined
 		let cachedContent: string | undefined = undefined
 		let cacheWriteTokens: number | undefined = undefined
 
+		// The minimum input token count for context caching is 4,096.
+		// For a basic appoximation we assume 4 characters per token.
path: src/api/providers/gemini.ts
user: login: ellipsis-dev[bot]
body: Typo in comments: 'appoximation' should be 'approximation' and 'accurat' should be 'accurate'.
```suggestion
		// For a basic approximation we assume 4 characters per token.
```

created_at: 2025-04-24T21:02:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,36 +48,72 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
+
 		let uncachedContent: Content[] | undefined = undefined
 		let cachedContent: string | undefined = undefined
 		let cacheWriteTokens: number | undefined = undefined
 
+		// The minimum input token count for context caching is 4,096.
+		// For a basic appoximation we assume 4 characters per token.
+		// We can use tiktoken eventually to get a more accurat token count.
 		// https://ai.google.dev/gemini-api/docs/caching?lang=node
-		// if (info.supportsPromptCache && cacheKey) {
-		// 	const cacheEntry = this.contentCaches.get(cacheKey)
+		// https://ai.google.dev/gemini-api/docs/tokens?lang=node
+		const isCacheAvailable =
+			info.supportsPromptCache &&
+			this.options.promptCachingEnabled &&
+			cacheKey &&
+			contentsLength > 4 * CONTEXT_CACHE_TOKEN_MINIMUM
path: src/api/providers/gemini.ts
user: login: ellipsis-dev[bot]
body: The condition uses a hardcoded multiplier (4) with `CONTEXT_CACHE_TOKEN_MINIMUM`. Consider extracting this multiplier into a named constant (e.g., `AVG_CHARS_PER_TOKEN`) for clarity and easier future adjustments.

<sup>This comment was generated because it violated a code review rule: [mrule_aQsEnH8jWdOfHq2Z](https://app.ellipsis.dev/RooVetGit/code-review/rules?id=mrule_aQsEnH8jWdOfHq2Z).</sup>
created_at: 2025-04-24T21:02:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,36 +50,76 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, thinkingConfig, maxOutputTokens, info } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
+		const contentsLength = systemInstruction.length + getMessagesLength(contents)
+
 		let uncachedContent: Content[] | undefined = undefined
 		let cachedContent: string | undefined = undefined
 		let cacheWriteTokens: number | undefined = undefined
 
+		// The minimum input token count for context caching is 4,096.
+		// For a basic appoximation we assume 4 characters per token.
+		// We can use tiktoken eventually to get a more accurat token count.
 		// https://ai.google.dev/gemini-api/docs/caching?lang=node
-		// if (info.supportsPromptCache && cacheKey) {
-		// 	const cacheEntry = this.contentCaches.get(cacheKey)
-
-		// 	if (cacheEntry) {
-		// 		uncachedContent = contents.slice(cacheEntry.count, contents.length)
-		// 		cachedContent = cacheEntry.key
-		// 	}
+		// https://ai.google.dev/gemini-api/docs/tokens?lang=node
+		const isCacheAvailable =
+			info.supportsPromptCache &&
+			this.options.promptCachingEnabled &&
+			cacheKey &&
+			contentsLength > 4 * CONTEXT_CACHE_TOKEN_MINIMUM
+
+		if (isCacheAvailable) {
+			const cacheEntry = this.contentCaches.get<CacheEntry>(cacheKey)
+
+			if (cacheEntry) {
+				uncachedContent = contents.slice(cacheEntry.count, contents.length)
+				cachedContent = cacheEntry.key
+				console.log(
+					`[GeminiHandler] using ${cacheEntry.count} cached messages (${cacheEntry.key}) and ${uncachedContent.length} uncached messages`,
+				)
+			}
 
-		// 	const newCacheEntry = await this.client.caches.create({
-		// 		model,
-		// 		config: { contents, systemInstruction, ttl: `${CACHE_TTL * 60}s` },
-		// 	})
+			if (!this.isCacheBusy) {
+				this.isCacheBusy = true
+				const timestamp = Date.now()
+
+				this.client.caches
+					.create({
+						model,
+						config: {
+							contents,
+							systemInstruction,
+							ttl: `${CACHE_TTL * 60}s`,
+							httpOptions: { timeout: 120_000 },
+						},
+					})
+					.then((result) => {
+						const { name, usageMetadata } = result
+
+						if (name) {
+							this.contentCaches.set<CacheEntry>(cacheKey, { key: name, count: contents.length })
+							cacheWriteTokens = usageMetadata?.totalTokenCount ?? 0
+							console.log(
path: src/api/providers/gemini.ts
user: login: ellipsis-dev[bot]
body: Consider using a centralized structured logging mechanism instead of direct `console.log` and `console.error` calls to adhere to our logging practices.

<sup>This comment was generated because it violated a code review rule: [mrule_OR1S8PRRHcvbdFib](https://app.ellipsis.dev/RooVetGit/code-review/rules?id=mrule_OR1S8PRRHcvbdFib).</sup>
created_at: 2025-04-24T21:21:12Z
author_association: CONTRIBUTOR

PR comments:
user: login: changeset-bot[bot]
created_at: 2025-04-24T20:13:02Z
body: ###  🦋  Changeset detected

Latest commit: c25c2f16e6e8d32d8f2ac3775e9a52062534de1f

**The changes in this PR will be included in the next version bump.**

<details><summary>This PR includes changesets to release 1 package</summary>

  | Name      | Type  |
| --------- | ----- |
| roo-cline | Patch |

</details>

Not sure what this means? [Click here  to learn what changesets are](https://github.com/changesets/changesets/blob/main/docs/adding-a-changeset.md).

[Click here if you're a maintainer who wants to add another changeset to this PR](https://github.com/RooVetGit/Roo-Code/new/cte/enable-gemini-prompt-caching?filename=.changeset/five-owls-pay.md&value=---%0A%22roo-cline%22%3A%20patch%0A---%0A%0AGemini%20caching%20improvements%0A)


author_association: NONE
user: login: ellipsis-dev[bot]
created_at: 2025-04-24T20:13:38Z
body: The changes in this pull request seem to be cohesive and related to improving the caching mechanism for the GeminiHandler. Therefore, it doesn't seem necessary to split this pull request into smaller ones. If you have any specific concerns or questions, feel free to reach out!
author_association: CONTRIBUTOR