diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index e778524c26..9f27915688 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -102,6 +102,136 @@ describe("GeminiHandler", () => { } }).rejects.toThrow() }) + + it("should filter out repetitive text chunks", async () => { + // Setup the mock to return repetitive text + ;(handler["client"].models.generateContentStream as any).mockResolvedValue({ + [Symbol.asyncIterator]: async function* () { + // Simulate the repetitive output from the bug report + yield { text: "I will now run the test." } + yield { text: "I will now run the test." } + yield { text: "I will now run the test." } + yield { text: "I will now run the test." } + yield { text: "I will now run the test." } + yield { text: " Different text." } + yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } } + }, + }) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + // Should have filtered out repetitive chunks + // First 3 occurrences should be kept (threshold is 3), rest filtered + const textChunks = chunks.filter((c) => c.type === "text") + expect(textChunks.length).toBe(4) // 3 repetitions + 1 different text + expect(textChunks[0]).toEqual({ type: "text", text: "I will now run the test." }) + expect(textChunks[1]).toEqual({ type: "text", text: "I will now run the test." }) + expect(textChunks[2]).toEqual({ type: "text", text: "I will now run the test." }) + expect(textChunks[3]).toEqual({ type: "text", text: " Different text." }) + }) + + it("should filter repetitive chunks from candidates structure", async () => { + // Setup mock with candidates structure containing repetitive text + ;(handler["client"].models.generateContentStream as any).mockResolvedValue({ + [Symbol.asyncIterator]: async function* () { + yield { + candidates: [ + { + content: { + parts: [{ text: "Repeating message" }], + }, + }, + ], + } + yield { + candidates: [ + { + content: { + parts: [{ text: "Repeating message" }], + }, + }, + ], + } + yield { + candidates: [ + { + content: { + parts: [{ text: "Repeating message" }], + }, + }, + ], + } + yield { + candidates: [ + { + content: { + parts: [{ text: "Repeating message" }], + }, + }, + ], + } + yield { + candidates: [ + { + content: { + parts: [{ text: "New content" }], + }, + }, + ], + } + yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } } + }, + }) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + // Should filter repetitions beyond threshold + const textChunks = chunks.filter((c) => c.type === "text") + expect(textChunks.length).toBe(4) // 3 repetitions (threshold) + 1 new content + expect(textChunks[0]).toEqual({ type: "text", text: "Repeating message" }) + expect(textChunks[3]).toEqual({ type: "text", text: "New content" }) + }) + + it("should reset repetition buffer for each new message", async () => { + // First message with repetitive text + ;(handler["client"].models.generateContentStream as any).mockResolvedValue({ + [Symbol.asyncIterator]: async function* () { + yield { text: "Test" } + yield { text: "Test" } + yield { text: "Test" } + yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } } + }, + }) + + const stream1 = handler.createMessage(systemPrompt, mockMessages) + const chunks1 = [] + for await (const chunk of stream1) { + chunks1.push(chunk) + } + + // Second message with same text should not be filtered (buffer reset) + const stream2 = handler.createMessage(systemPrompt, mockMessages) + const chunks2 = [] + for await (const chunk of stream2) { + chunks2.push(chunk) + } + + // Both messages should have the same chunks (buffer was reset) + const textChunks1 = chunks1.filter((c) => c.type === "text") + const textChunks2 = chunks2.filter((c) => c.type === "text") + expect(textChunks1.length).toBe(3) + expect(textChunks2.length).toBe(3) + }) }) describe("completePrompt", () => { diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index 45be6977d9..7c59aa9772 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -32,6 +32,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl private client: GoogleGenAI private lastThoughtSignature?: string private lastResponseId?: string + private recentTextChunks: string[] = [] + private readonly maxRecentChunks = 5 + private readonly repetitionThreshold = 3 constructor({ isVertex, ...options }: GeminiHandlerOptions) { super() @@ -72,6 +75,8 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl // Reset per-request metadata that we persist into apiConversationHistory. this.lastThoughtSignature = undefined this.lastResponseId = undefined + // Reset repetition detection buffer for each new message + this.recentTextChunks = [] // For hybrid/budget reasoning models (e.g. Gemini 2.5 Pro), respect user-configured // modelMaxTokens so the ThinkingBudget slider can control the cap. For effort-only or @@ -263,7 +268,15 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } else { // This is regular content if (part.text) { - yield { type: "text", text: part.text } + // Check for repetition before yielding + if (!this.isRepetitiveChunk(part.text)) { + yield { type: "text", text: part.text } + this.addToRecentChunks(part.text) + } else { + console.warn( + `Gemini: Filtered repetitive chunk: "${part.text.substring(0, 50)}..."`, + ) + } } } } @@ -272,7 +285,13 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl // Fallback to the original text property if no candidates structure else if (chunk.text) { - yield { type: "text", text: chunk.text } + // Check for repetition before yielding + if (!this.isRepetitiveChunk(chunk.text)) { + yield { type: "text", text: chunk.text } + this.addToRecentChunks(chunk.text) + } else { + console.warn(`Gemini: Filtered repetitive chunk: "${chunk.text.substring(0, 50)}..."`) + } } if (chunk.usageMetadata) { @@ -501,4 +520,31 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return totalCost } + + /** + * Check if a text chunk is repetitive based on recent chunks + */ + private isRepetitiveChunk(text: string): boolean { + if (!text || text.trim().length === 0) { + return false + } + + // Count how many times this exact text appears in recent chunks + const occurrences = this.recentTextChunks.filter((chunk) => chunk === text).length + + // If the text appears more than the threshold, it's considered repetitive + return occurrences >= this.repetitionThreshold + } + + /** + * Add a text chunk to the recent chunks buffer + */ + private addToRecentChunks(text: string): void { + this.recentTextChunks.push(text) + + // Keep only the most recent chunks + if (this.recentTextChunks.length > this.maxRecentChunks) { + this.recentTextChunks.shift() + } + } }