Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand Down Expand Up @@ -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", () => {
Expand Down
15 changes: 5 additions & 10 deletions src/core/assistant-message/presentAssistantMessage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading