Unify effective tool policy across prompts and tool filtering

apps/vscode-e2e/src/suite/restart-persistence.test.ts

@@ -25,6 +25,19 @@ async function quitGracefully(): Promise<void> {
 	await vscode.commands.executeCommand("workbench.action.quit")
 }
 
+// The API history file is committed via a backup-rename swap under an advisory
+// lock, so an immediate post-completion read can transiently observe it as
+// missing. Poll the same read the assertion uses until it reports a non-empty
+// history; the assertion semantics below are unchanged.
+async function waitForApiConversationHistoryLength(api: RooCodeAPI, taskId: string): Promise<number> {
+	let length = 0
+	await waitFor(async () => {
+		length = await api.getTaskApiConversationHistoryLength(taskId)
+		return length > 0
+	})
+	return length
+}
+
 async function runCreate(api: RooCodeAPI): Promise<void> {
 	let taskId: string | undefined
 	let createPhasePassed = false
@@ -43,6 +56,11 @@ async function runCreate(api: RooCodeAPI): Promise<void> {
 		})
 		await waitUntilCompleted({ api, taskId })
 		assert.strictEqual(sawMarker, true, `Completion should include ${MARKER}`)
+		const historyItem = await api.getTaskHistoryItem(taskId)
+		assert.ok(historyItem, "Completed task should have a history item")
+		assert.ok(historyItem.task.includes("RESTART_PERSISTENCE_SMOKE"), "History title should include the marker")
+		const conversationLength = await waitForApiConversationHistoryLength(api, taskId)
+		assert.ok(conversationLength > 0, "Completed task should persist API conversation history")
 
 		const result: PhaseResult = {
 			version: PHASE_RESULT_VERSION,

src/api/index.ts

@@ -130,8 +130,15 @@ export interface ApiHandler {
 	 * Ensures model metadata has been fetched from the remote API so that getModel()
 	 * returns accurate info (context window, pricing, etc.) instead of hardcoded defaults.
 	 * Only router providers that discover models over the network implement this.
+	 *
+	 * `signal` bounds the caller's wait: when it aborts (e.g. the caller's bounded
+	 * metadata wait expired or the owning task was cancelled), the returned promise
+	 * settles with a rejection so no handler-side waiter outlives its caller.
+	 * Fetchers that observe the signal may also stop their network request; the
+	 * shared, de-duplicated catalog fetch may still complete and populate the model
+	 * cache, which is by design for concurrent waiters.
 	 */
-	ensureModelFetched?(): Promise<void>
+	ensureModelFetched?(signal?: AbortSignal): Promise<void>
 
 	/**
 	 * Optional context window for context-management / auto-condense when it must differ from

src/api/providers/__tests__/zoo-gateway.spec.ts

@@ -724,6 +724,72 @@ describe("ZooGatewayHandler", () => {
 			expect(refreshModels).not.toHaveBeenCalled()
 		})
 
+		it("settles the waiter with a rejection when the signal aborts mid-fetch", async () => {
+			// A caller that gives up must not leave a handler-side waiter pending
+			// on the (shared) catalog fetch: with an observing signal, the
+			// ensureModelFetched promise rejects at abort time, while the
+			// underlying fetch continues untouched for any other waiter.
+			const { getModels } = await import("../fetchers/modelCache")
+			vitest.mocked(getModels).mockImplementationOnce(() => new Promise(() => {}))
+
+			const handler = new ZooGatewayHandler(mockOptions)
+			const controller = new AbortController()
+
+			const wait = handler.ensureModelFetched(controller.signal)
+			// Let the waiter attach its abort listener before cancelling.
+			await Promise.resolve()
+			controller.abort()
+
+			await expect(wait).rejects.toThrow()
+		})
+
+		it("settles the waiter when the fetch wins against a live signal and detaches the listener", async () => {
+			// Fetch-wins branch: resolve() must settle the await (a dropped
+			// resolve or a detached .then handler hangs this test), the abort
+			// listener must be registered with the real { once: true } options
+			// object, and the detach must target the *same* event name/handler
+			// pair that was registered — a mutated event name detaches nothing.
+			const handler = new ZooGatewayHandler(mockOptions)
+			const controller = new AbortController()
+			const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener")
+			const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener")
+
+			await handler.ensureModelFetched(controller.signal)
+
+			expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true })
+			const registered = addEventListenerSpy.mock.calls.find(([event]) => event === "abort")
+			expect(registered).toBeDefined()
+			expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registered?.[1])
+		})
+
+		it("rejects a signal-observing waiter with the fetch error and detaches the listener", async () => {
+			// Rejection-branch twin of the fetch-wins test: reject(error) must
+			// propagate the catalog failure to the waiter (a dropped reject hangs
+			// this test) and the listener must be detached under the right event
+			// name. Existing reject coverage passes no signal, so this branch was
+			// never executed before.
+			const { getModels } = await import("../fetchers/modelCache")
+			vitest.mocked(getModels).mockRejectedValueOnce(new Error("network down"))
+
+			const handler = new ZooGatewayHandler(mockOptions)
+			const controller = new AbortController()
+			const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener")
+			const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener")
+
+			await expect(handler.ensureModelFetched(controller.signal)).rejects.toThrow("network down")
+			const registered = addEventListenerSpy.mock.calls.find(([event]) => event === "abort")
+			expect(registered).toBeDefined()
+			expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registered?.[1])
+		})
+
+		it("never starts a wait when the signal is already aborted", async () => {
+			const handler = new ZooGatewayHandler(mockOptions)
+			const controller = new AbortController()
+			controller.abort()
+
+			await expect(handler.ensureModelFetched(controller.signal)).rejects.toThrow()
+		})
+
 		it("skips the fetch when models are already populated", async () => {
 			const handler = new ZooGatewayHandler(mockOptions)
 			const { getModels, refreshModels } = await import("../fetchers/modelCache")

src/api/providers/router-provider.ts

@@ -108,8 +108,35 @@ export abstract class RouterProvider extends BaseProvider {
 		return this.modelFetchPromise
 	}
 
-	async ensureModelFetched(): Promise<void> {
-		await this.fetchModel()
+	async ensureModelFetched(signal?: AbortSignal): Promise<void> {
+		// A caller that already gave up must not start (or keep) a wait on the
+		// shared catalog fetch.
+		if (signal?.aborted) {
+			throw signal.reason
+		}
+
+		const fetch = this.fetchModel()
+		if (!signal) {
+			await fetch
+			return
+		}
+
+		// Detach this waiter as soon as the signal aborts; the shared in-flight
+		// fetch continues for any other waiter and still populates the cache.
+		await new Promise<void>((resolve, reject) => {
+			const onAbort = () => reject(signal.reason)
+			signal.addEventListener("abort", onAbort, { once: true })
+			fetch.then(
+				() => {
+					signal.removeEventListener("abort", onAbort)
+					resolve()
+				},
+				(error: unknown) => {
+					signal.removeEventListener("abort", onAbort)
+					reject(error)
+				},
+			)
+		})
 	}
 
 	override getModel(): { id: string; info: ModelInfo } {

src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts

@@ -23,6 +23,15 @@ vi.mock("@roo-code/core", () => ({
 	},
 }))
 
+// Mock the tool handlers so the tests only exercise validation (toolRequirements)
+// and never the real tool execution logic.
+vi.mock("../../tools/AttemptCompletionTool", () => ({
+	attemptCompletionTool: { handle: vi.fn().mockResolvedValue(undefined) },
+}))
+vi.mock("../../tools/AskFollowupQuestionTool", () => ({
+	askFollowupQuestionTool: { handle: vi.fn().mockResolvedValue(undefined) },
+}))
+
 // presentAssistantMessage records tool usage through TelemetryService.instance.
 vi.mock("@roo-code/telemetry", () => ({
 	TelemetryService: {
@@ -333,6 +342,86 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
 				edit: false,
 			})
 		})
+
+		it("never marks a protocol tool (attempt_completion) as blocked", async () => {
+			mockTask.assistantMessageContent = [
+				{
+					type: "tool_use",
+					id: "tool_call_protocol_123",
+					name: "attempt_completion",
+					params: {},
+					nativeArgs: {},
+					partial: false,
+				},
+			]
+
+			mockTask.providerRef = {
+				deref: () => ({
+					getState: vi.fn().mockResolvedValue({
+						mode: "code",
+						customModes: [],
+						experiments: {
+							customTools: false,
+						},
+						disabledTools: ["attempt_completion"],
+					}),
+				}),
+			}
+
+			await presentAssistantMessage(mockTask)
+
+			const validateToolUseMock = vi.mocked(validateToolUse)
+			expect(validateToolUseMock).toHaveBeenCalled()
+			const toolRequirements = validateToolUseMock.mock.calls[0][3]
+			// Protocol tools never enter toolRequirements, so the validator cannot
+			// block them even when disabledTools lists them.
+			expect(toolRequirements).not.toHaveProperty("attempt_completion")
+
+			// With validateToolUse mocked to return normally, the block proceeds
+			// past validation: no validation-error tool_result is pushed.
+			const errorToolResults = mockTask.userMessageContent.filter((block: unknown) => {
+				const b = block as { type?: string; is_error?: boolean }
+				return b.type === "tool_result" && b.is_error
+			})
+			expect(errorToolResults).toEqual([])
+		})
+
+		it("still marks ordinary tools (ask_followup_question) as blocked", async () => {
+			mockTask.assistantMessageContent = [
+				{
+					type: "tool_use",
+					id: "tool_call_ordinary_123",
+					name: "ask_followup_question",
+					params: { question: "Which option?" },
+					nativeArgs: { question: "Which option?" },
+					partial: false,
+				},
+			]
+
+			mockTask.providerRef = {
+				deref: () => ({
+					getState: vi.fn().mockResolvedValue({
+						mode: "code",
+						customModes: [],
+						experiments: {
+							customTools: false,
+						},
+						disabledTools: ["ask_followup_question"],
+					}),
+				}),
+			}
+
+			await presentAssistantMessage(mockTask)
+
+			const validateToolUseMock = vi.mocked(validateToolUse)
+			expect(validateToolUseMock).toHaveBeenCalled()
+			const toolRequirements = validateToolUseMock.mock.calls[0][3]
+			// Control/ordinary tools remain blockable — the inverse of the
+			// protocol-tool guarantee.
+			expect(toolRequirements).toMatchObject({
+				ask_followup_question: false,
+			})
+		})
 	})
 
 	describe("Partial blocks", () => {

src/core/assistant-message/presentAssistantMessage.ts

@@ -36,6 +36,7 @@ import { skillTool } from "../tools/SkillTool"
 import { generateImageTool } from "../tools/GenerateImageTool"
 import { applyDiffTool as applyDiffToolClass } from "../tools/ApplyDiffTool"
 import { isValidToolName, validateToolUse } from "../tools/validateToolUse"
+import { buildToolRequirements } from "../prompts/tools/effective-tool-policy"
 import { codebaseSearchTool } from "../tools/CodebaseSearchTool"
 
 import { formatResponse } from "../prompts/responses"
@@ -604,16 +605,10 @@ export async function presentAssistantMessage(cline: Task) {
 				const isCustomTool = Boolean(stateExperiments?.customTools && customToolRegistry.has(block.name))
 
 				try {
-					const toolRequirements =
-						disabledTools?.reduce(
-							(acc: Record<string, boolean>, tool: string) => {
-								acc[tool] = false
-								const resolvedToolName = resolveToolAlias(tool)
-								acc[resolvedToolName] = false
-								return acc
-							},
-							{} as Record<string, boolean>,
-						) ?? {}
+					// Use the exported resolver so `attempt_completion` (and its aliases)
+					// never enters `toolRequirements` — the runtime validator never blocks a
+					// protocol tool. See `buildToolRequirements` in effective-tool-policy.ts.
+					const toolRequirements = buildToolRequirements(disabledTools)
 
 					validateToolUse(
 						block.name as ToolName,