fix(memory): rate-limit Gemini embedding requests and cache per-batch#42578
fix(memory): rate-limit Gemini embedding requests and cache per-batch#42578chencheng-li wants to merge 2 commits into
Conversation
Greptile SummaryThis PR introduces a Promise-chain rate limiter in Key finding: The rate-limit queue implementation has a scope issue: The per-batch cache improvement is solid and addresses the data loss risk from interruptions during safe reindex operations. Confidence Score: 2/5
Last reviewed commit: 78e6d4f |
| // Global rate limiter: serialize all Gemini embedding requests across | ||
| // concurrent callers. Uses batchEmbedContents to pack multiple texts | ||
| // into one request (up to 8192 token limit), then waits before the | ||
| // next request to stay under free-tier 100 RPM. | ||
| let rateLimitQueue: Promise<void> = Promise.resolve(); | ||
| const GEMINI_EMBED_DELAY_MS = 1200; | ||
|
|
||
| const rateLimitedBatch = async (texts: string[]): Promise<number[][]> => { | ||
| return new Promise((resolve, reject) => { | ||
| rateLimitQueue = rateLimitQueue.then(async () => { | ||
| try { | ||
| const requests = texts.map((text) => ({ | ||
| model: client.modelPath, | ||
| content: { parts: [{ text }] }, | ||
| taskType: "RETRIEVAL_DOCUMENT", | ||
| })); | ||
| debugEmbeddingsLog("gemini rate-limited batch: sending", { | ||
| texts: texts.length, | ||
| totalChars: texts.reduce((sum, t) => sum + t.length, 0), | ||
| }); | ||
| const payload = await executeWithApiKeyRotation({ | ||
| provider: "google", | ||
| apiKeys: client.apiKeys, | ||
| execute: (apiKey) => fetchWithGeminiAuth(apiKey, batchUrl, { requests }), | ||
| }); | ||
| const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : []; | ||
| debugEmbeddingsLog("gemini rate-limited batch: success", { | ||
| texts: texts.length, | ||
| embeddingsReturned: embeddings.length, | ||
| dims: embeddings[0]?.values?.length ?? 0, | ||
| }); | ||
| resolve(texts.map((_, index) => embeddings[index]?.values ?? [])); | ||
| } catch (err) { | ||
| debugEmbeddingsLog("gemini rate-limited batch: failed", { | ||
| texts: texts.length, | ||
| error: err instanceof Error ? err.message : String(err), | ||
| }); | ||
| reject(err); | ||
| } | ||
| await new Promise((r) => setTimeout(r, GEMINI_EMBED_DELAY_MS)); | ||
| }); | ||
| }); | ||
| }; |
There was a problem hiding this comment.
Rate-limit queue is per-instance, not truly global
The comment says "Global rate limiter: serialize all Gemini embedding requests across concurrent callers," but rateLimitQueue is a closure variable created inside createGeminiEmbeddingProvider. Each call to createGeminiEmbeddingProvider produces an independent queue. If the provider is instantiated more than once (e.g., multiple agents each creating their own memory manager), the queues are completely independent and concurrent requests from different instances will not be serialized, which can still exceed the 100 RPM free-tier limit.
To make the queue truly global, move it to module scope:
| // Global rate limiter: serialize all Gemini embedding requests across | |
| // concurrent callers. Uses batchEmbedContents to pack multiple texts | |
| // into one request (up to 8192 token limit), then waits before the | |
| // next request to stay under free-tier 100 RPM. | |
| let rateLimitQueue: Promise<void> = Promise.resolve(); | |
| const GEMINI_EMBED_DELAY_MS = 1200; | |
| const rateLimitedBatch = async (texts: string[]): Promise<number[][]> => { | |
| return new Promise((resolve, reject) => { | |
| rateLimitQueue = rateLimitQueue.then(async () => { | |
| try { | |
| const requests = texts.map((text) => ({ | |
| model: client.modelPath, | |
| content: { parts: [{ text }] }, | |
| taskType: "RETRIEVAL_DOCUMENT", | |
| })); | |
| debugEmbeddingsLog("gemini rate-limited batch: sending", { | |
| texts: texts.length, | |
| totalChars: texts.reduce((sum, t) => sum + t.length, 0), | |
| }); | |
| const payload = await executeWithApiKeyRotation({ | |
| provider: "google", | |
| apiKeys: client.apiKeys, | |
| execute: (apiKey) => fetchWithGeminiAuth(apiKey, batchUrl, { requests }), | |
| }); | |
| const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : []; | |
| debugEmbeddingsLog("gemini rate-limited batch: success", { | |
| texts: texts.length, | |
| embeddingsReturned: embeddings.length, | |
| dims: embeddings[0]?.values?.length ?? 0, | |
| }); | |
| resolve(texts.map((_, index) => embeddings[index]?.values ?? [])); | |
| } catch (err) { | |
| debugEmbeddingsLog("gemini rate-limited batch: failed", { | |
| texts: texts.length, | |
| error: err instanceof Error ? err.message : String(err), | |
| }); | |
| reject(err); | |
| } | |
| await new Promise((r) => setTimeout(r, GEMINI_EMBED_DELAY_MS)); | |
| }); | |
| }); | |
| }; | |
| // module-level, shared across all provider instances | |
| let _globalGeminiRateLimitQueue: Promise<void> = Promise.resolve(); | |
| const GEMINI_EMBED_DELAY_MS = 1200; | |
| export async function createGeminiEmbeddingProvider(...) { | |
| // ... | |
| const rateLimitedBatch = async (texts: string[]): Promise<number[][]> => { | |
| return new Promise((resolve, reject) => { | |
| _globalGeminiRateLimitQueue = _globalGeminiRateLimitQueue.then(async () => { | |
| // ... same body ... | |
| await new Promise((r) => setTimeout(r, GEMINI_EMBED_DELAY_MS)); | |
| }); | |
| }); | |
| }; | |
| // ... | |
| } |
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/memory/embeddings-gemini.ts
Line: 121-163
Comment:
**Rate-limit queue is per-instance, not truly global**
The comment says "Global rate limiter: serialize all Gemini embedding requests across concurrent callers," but `rateLimitQueue` is a closure variable created inside `createGeminiEmbeddingProvider`. Each call to `createGeminiEmbeddingProvider` produces an independent queue. If the provider is instantiated more than once (e.g., multiple agents each creating their own memory manager), the queues are completely independent and concurrent requests from different instances will not be serialized, which can still exceed the 100 RPM free-tier limit.
To make the queue truly global, move it to module scope:
```suggestion
// module-level, shared across all provider instances
let _globalGeminiRateLimitQueue: Promise<void> = Promise.resolve();
const GEMINI_EMBED_DELAY_MS = 1200;
export async function createGeminiEmbeddingProvider(...) {
// ...
const rateLimitedBatch = async (texts: string[]): Promise<number[][]> => {
return new Promise((resolve, reject) => {
_globalGeminiRateLimitQueue = _globalGeminiRateLimitQueue.then(async () => {
// ... same body ...
await new Promise((r) => setTimeout(r, GEMINI_EMBED_DELAY_MS));
});
});
};
// ...
}
```
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78e6d4ff51
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| rateLimitQueue = rateLimitQueue.then(async () => { | ||
| try { |
There was a problem hiding this comment.
Add cancellation path for queued Gemini batch jobs
The new rateLimitQueue serializes every embedBatch call behind the previous callback, but the queued callback awaits the HTTP request without any abort signal. embedBatchWithRetry times out callers after 120s, yet that timeout does not cancel the in-flight queued job, so one stuck Gemini request can block the queue and cause all later batches to time out/retry without ever being sent. This regression only appears after introducing the shared queue and can stall memory index during network hangs or very slow Gemini responses.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9a4f2c7118
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| geminiEmbedRateLimitQueue = geminiEmbedRateLimitQueue.then(async () => { | ||
| try { |
There was a problem hiding this comment.
Cancel timed-out Gemini jobs in the shared queue
Fresh evidence: this commit introduces the module-level geminiEmbedRateLimitQueue here, while embedBatchWithRetry still enforces timeout via withTimeout (src/memory/manager-embedding-ops.ts:519) without aborting the underlying HTTP work. If one queued Gemini request hangs, the caller times out after 120s but the queued callback keeps running, so later embedBatch calls remain blocked behind it and can time out before their requests are even sent, causing cascading failures/stalls during memory index under slow or stuck network responses.
Useful? React with 👍 / 👎.
deb33ef to
6753d3e
Compare
|
Closing and replacing with a clean follow-up PR rebased on latest main. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb5a462e61
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const payload = await executeWithApiKeyRotation({ | ||
| provider: "google", | ||
| apiKeys: client.apiKeys, | ||
| execute: (apiKey) => fetchWithGeminiAuth(apiKey, batchUrl, { requests }), |
There was a problem hiding this comment.
Pass abort signal into queued Gemini batch request
rateLimitedBatch creates an AbortController and a 30s timeout, but the queued request path never passes controller.signal into fetchWithGeminiAuth. In a slow or hung Gemini call, the timeout callback runs without cancelling the HTTP work, so that queue slot can remain occupied until lower-level network timeouts while later batches pile up behind it and hit caller-level timeouts/retries. Threading the signal through fetchWithGeminiAuth and withRemoteHttpResponse is needed for the per-item timeout to actually protect queue progress.
Useful? React with 👍 / 👎.
Problem
When using Gemini free-tier embedding API (
gemini-embedding-2-preview),openclaw memory indexfails with 429 errors because:embedBatchin parallel, exceeding the 100 RPM limitSolution
embeddings-gemini.ts: serializes all Gemini embedding requests through a Promise queue with 1.2s delay between requests (~50 RPM, well under 100 RPM limit)manager-embedding-ops.ts:upsertEmbeddingCacheruns after each successful batch, not at the endTesting
pnpm check✅vitest run src/memory/embeddings.test.ts— 23/23 tests pass ✅🤖 AI-assisted (OpenClaw + Claude Opus)