Skip to content

Commit df215b1

Browse files
roomoteroomote
committed
fix: prevent Gemini grounding loop in Ask mode
- Add detection for Gemini grounding responses with citations - Prevent consecutive mistake count increment when grounding is used - Fix infinite loop issue when Gemini provides search results - Add comprehensive test coverage for grounding detection Fixes #6503
1 parent 74672fa commit df215b1

File tree

2 files changed

+263
-1
lines changed

2 files changed

+263
-1
lines changed

‎src/core/task/Task.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1737,7 +1737,21 @@ export class Task extends EventEmitter<TaskEvents> {
17371737
// either use a tool or attempt_completion.
17381738
const didToolUse = this.assistantMessageContent.some((block) => block.type === "tool_use")
17391739

1740-
if (!didToolUse) {
1740+
// Check if this is a Gemini grounding response (contains citations/sources)
1741+
// Gemini grounding responses include citations that indicate tool usage (search)
1742+
const hasGeminiGrounding =
1743+
this.apiConfiguration.apiProvider === "gemini" &&
1744+
this.assistantMessageContent.some(
1745+
(block) =>
1746+
block.type === "text" &&
1747+
block.content &&
1748+
(block.content.includes("Sources:") ||
1749+
block.content.includes("[1]") ||
1750+
block.content.includes("[2]") ||
1751+
/\[\d+\]\(https?:\/\/[^\)]+\)/.test(block.content)),
1752+
)
1753+
1754+
if (!didToolUse && !hasGeminiGrounding) {
17411755
this.userMessageContent.push({ type: "text", text: formatResponse.noToolsUsed() })
17421756
this.consecutiveMistakeCount++
17431757
}

‎src/core/task/__tests__/Task.gemini-grounding.spec.ts‎

Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
// npx vitest run src/core/task/__tests__/Task.gemini-grounding.spec.ts
2+
3+
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
4+
import { Task } from "../Task"
5+
import type { ClineProvider } from "../../webview/ClineProvider"
6+
import type { ProviderSettings } from "@roo-code/types"
7+
8+
describe("Task Gemini Grounding Loop Prevention", () => {
9+
let mockProvider: Partial<ClineProvider>
10+
let mockApiConfiguration: ProviderSettings
11+
12+
beforeEach(() => {
13+
mockProvider = {
14+
context: {
15+
globalStorageUri: { fsPath: "/tmp/test" },
16+
} as any,
17+
getState: vi.fn().mockResolvedValue({
18+
mode: "ask",
19+
customModes: [],
20+
experiments: {},
21+
}),
22+
postStateToWebview: vi.fn(),
23+
updateTaskHistory: vi.fn(),
24+
log: vi.fn(),
25+
}
26+
27+
mockApiConfiguration = {
28+
apiProvider: "gemini",
29+
apiModelId: "gemini-2.5-pro",
30+
enableGrounding: true,
31+
} as ProviderSettings
32+
})
33+
34+
afterEach(() => {
35+
vi.clearAllMocks()
36+
})
37+
38+
it("should not increment mistake count when Gemini provides grounding citations", async () => {
39+
const task = new Task({
40+
provider: mockProvider as ClineProvider,
41+
apiConfiguration: mockApiConfiguration,
42+
task: "Test task",
43+
startTask: false,
44+
})
45+
46+
// Wait for task mode initialization
47+
await task.waitForModeInitialization()
48+
49+
// Simulate Gemini response with grounding citations
50+
task.assistantMessageContent = [
51+
{
52+
type: "text",
53+
content:
54+
"Here's information about Tailwind CSS v4:\n\nTailwind CSS v4 introduces several new features...\n\nSources: [1](https://tailwindcss.com/docs), [2](https://github.com/tailwindlabs/tailwindcss)",
55+
partial: false,
56+
},
57+
]
58+
59+
// Set up the state that would normally be set during streaming
60+
task.userMessageContentReady = true
61+
task.userMessageContent = []
62+
63+
const initialMistakeCount = task.consecutiveMistakeCount
64+
65+
// Simulate the tool usage detection logic
66+
const didToolUse = task.assistantMessageContent.some((block) => block.type === "tool_use")
67+
68+
// Check if this is a Gemini grounding response (contains citations/sources)
69+
const hasGeminiGrounding =
70+
task.apiConfiguration.apiProvider === "gemini" &&
71+
task.assistantMessageContent.some(
72+
(block) =>
73+
block.type === "text" &&
74+
block.content &&
75+
(block.content.includes("Sources:") ||
76+
block.content.includes("[1]") ||
77+
block.content.includes("[2]") ||
78+
/\[\d+\]\(https?:\/\/[^\)]+\)/.test(block.content)),
79+
)
80+
81+
if (!didToolUse && !hasGeminiGrounding) {
82+
task.consecutiveMistakeCount++
83+
}
84+
85+
// Verify that mistake count was NOT incremented due to grounding detection
86+
expect(task.consecutiveMistakeCount).toBe(initialMistakeCount)
87+
expect(hasGeminiGrounding).toBe(true)
88+
expect(didToolUse).toBe(false)
89+
})
90+
91+
it("should increment mistake count when Gemini provides response without grounding or tools", async () => {
92+
const task = new Task({
93+
provider: mockProvider as ClineProvider,
94+
apiConfiguration: mockApiConfiguration,
95+
task: "Test task",
96+
startTask: false,
97+
})
98+
99+
// Wait for task mode initialization
100+
await task.waitForModeInitialization()
101+
102+
// Simulate Gemini response without grounding citations or tools
103+
task.assistantMessageContent = [
104+
{
105+
type: "text",
106+
content: "Here's some general information about web development without any sources.",
107+
partial: false,
108+
},
109+
]
110+
111+
// Set up the state that would normally be set during streaming
112+
task.userMessageContentReady = true
113+
task.userMessageContent = []
114+
115+
const initialMistakeCount = task.consecutiveMistakeCount
116+
117+
// Simulate the tool usage detection logic
118+
const didToolUse = task.assistantMessageContent.some((block) => block.type === "tool_use")
119+
120+
// Check if this is a Gemini grounding response (contains citations/sources)
121+
const hasGeminiGrounding =
122+
task.apiConfiguration.apiProvider === "gemini" &&
123+
task.assistantMessageContent.some(
124+
(block) =>
125+
block.type === "text" &&
126+
block.content &&
127+
(block.content.includes("Sources:") ||
128+
block.content.includes("[1]") ||
129+
block.content.includes("[2]") ||
130+
/\[\d+\]\(https?:\/\/[^\)]+\)/.test(block.content)),
131+
)
132+
133+
if (!didToolUse && !hasGeminiGrounding) {
134+
task.consecutiveMistakeCount++
135+
}
136+
137+
// Verify that mistake count WAS incremented since no grounding was detected
138+
expect(task.consecutiveMistakeCount).toBe(initialMistakeCount + 1)
139+
expect(hasGeminiGrounding).toBe(false)
140+
expect(didToolUse).toBe(false)
141+
})
142+
143+
it("should detect various grounding citation formats", async () => {
144+
const task = new Task({
145+
provider: mockProvider as ClineProvider,
146+
apiConfiguration: mockApiConfiguration,
147+
task: "Test task",
148+
startTask: false,
149+
})
150+
151+
// Wait for task mode initialization
152+
await task.waitForModeInitialization()
153+
154+
const testCases = [
155+
{
156+
content: "Information here.\n\nSources: [1](https://example.com)",
157+
shouldDetectGrounding: true,
158+
description: "Sources with numbered links",
159+
},
160+
{
161+
content: "Some info [1](https://example.com) and more [2](https://another.com)",
162+
shouldDetectGrounding: true,
163+
description: "Inline numbered citations",
164+
},
165+
{
166+
content: "Text with [1] reference",
167+
shouldDetectGrounding: true,
168+
description: "Simple numbered reference",
169+
},
170+
{
171+
content: "Just regular text without any citations",
172+
shouldDetectGrounding: false,
173+
description: "No citations",
174+
},
175+
{
176+
content: "Text with [abc] but not numbered",
177+
shouldDetectGrounding: false,
178+
description: "Non-numbered brackets",
179+
},
180+
]
181+
182+
for (const testCase of testCases) {
183+
task.assistantMessageContent = [
184+
{
185+
type: "text",
186+
content: testCase.content,
187+
partial: false,
188+
},
189+
]
190+
191+
const hasGeminiGrounding =
192+
task.apiConfiguration.apiProvider === "gemini" &&
193+
task.assistantMessageContent.some(
194+
(block) =>
195+
block.type === "text" &&
196+
block.content &&
197+
(block.content.includes("Sources:") ||
198+
block.content.includes("[1]") ||
199+
block.content.includes("[2]") ||
200+
/\[\d+\]\(https?:\/\/[^\)]+\)/.test(block.content)),
201+
)
202+
203+
expect(hasGeminiGrounding, `Failed for case: ${testCase.description}`).toBe(testCase.shouldDetectGrounding)
204+
}
205+
})
206+
207+
it("should only apply grounding detection for Gemini provider", async () => {
208+
// Test with non-Gemini provider
209+
const nonGeminiConfig: ProviderSettings = {
210+
apiProvider: "anthropic",
211+
apiModelId: "claude-3-5-sonnet-20241022",
212+
} as ProviderSettings
213+
214+
const task = new Task({
215+
provider: mockProvider as ClineProvider,
216+
apiConfiguration: nonGeminiConfig,
217+
task: "Test task",
218+
startTask: false,
219+
})
220+
221+
// Wait for task mode initialization
222+
await task.waitForModeInitialization()
223+
224+
// Simulate response with citation-like content from non-Gemini provider
225+
task.assistantMessageContent = [
226+
{
227+
type: "text",
228+
content: "Here's information.\n\nSources: [1](https://example.com)",
229+
partial: false,
230+
},
231+
]
232+
233+
const hasGeminiGrounding =
234+
task.apiConfiguration.apiProvider === "gemini" &&
235+
task.assistantMessageContent.some(
236+
(block) =>
237+
block.type === "text" &&
238+
block.content &&
239+
(block.content.includes("Sources:") ||
240+
block.content.includes("[1]") ||
241+
block.content.includes("[2]") ||
242+
/\[\d+\]\(https?:\/\/[^\)]+\)/.test(block.content)),
243+
)
244+
245+
// Should not detect grounding for non-Gemini providers
246+
expect(hasGeminiGrounding).toBe(false)
247+
})
248+
})

0 commit comments

Comments
 (0)