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
1 change: 1 addition & 0 deletions packages/types/src/tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ export const toolNames = [
"skill",
"generate_image",
"custom_tool",
"invalid_tool_call",
] as const

export const toolNamesSchema = z.enum(toolNames)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,310 @@
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts

import type { Anthropic } from "@anthropic-ai/sdk"
import { describe, it, expect, beforeEach, vi } from "vitest"
import { presentAssistantMessage } from "../presentAssistantMessage"
import { validateToolUse } from "../../tools/validateToolUse"
import { getModeBySlug } from "../../../shared/modes"
import type { Task } from "../../task/Task"

vi.mock("../../task/Task")
vi.mock("../../../shared/modes", async (importOriginal) => {
const actual = await importOriginal<typeof import("../../../shared/modes")>()
return {
...actual,
getModeBySlug: vi.fn(actual.getModeBySlug),
}
})
vi.mock("../../tools/validateToolUse", () => ({
validateToolUse: vi.fn(),
isValidToolName: vi.fn((toolName: string) =>
["read_file", "write_to_file", "ask_followup_question", "attempt_completion", "use_mcp_tool"].includes(
toolName,
),
),
}))

vi.mock("@roo-code/core", () => ({
customToolRegistry: {
has: vi.fn(() => false),
get: vi.fn(),
},
}))

vi.mock("@roo-code/telemetry", () => ({
TelemetryService: {
instance: {
captureToolUsage: vi.fn(),
captureConsecutiveMistakeError: vi.fn(),
captureEvent: vi.fn(),
},
},
}))

import { TelemetryService } from "@roo-code/telemetry"

interface MockTask {
taskId: string
instanceId: string
abort: boolean
presentAssistantMessageLocked: boolean
presentAssistantMessageHasPendingUpdates: boolean
currentStreamingContentIndex: number
assistantMessageContent: unknown[]
userMessageContent: Anthropic.ToolResultBlockParam[]
didCompleteReadingStream: boolean
didRejectTool: boolean
didAlreadyUseTool: boolean
consecutiveMistakeCount: number
clineMessages: unknown[]
api: { getModel: () => { id: string; info: Record<string, unknown> } }
recordToolUsage: ReturnType<typeof vi.fn>
recordToolError: ReturnType<typeof vi.fn>
toolRepetitionDetector: { check: ReturnType<typeof vi.fn> }
providerRef: {
deref: () => {
getState: ReturnType<typeof vi.fn>
getMcpHub?: () => { findServerNameBySanitizedName: (name: string) => string | undefined }
}
}
say: ReturnType<typeof vi.fn>
ask: ReturnType<typeof vi.fn>
pushToolResultToUserContent: ReturnType<typeof vi.fn>
}

describe("presentAssistantMessage - tool usage attribution", () => {
let mockTask: MockTask

beforeEach(() => {
vi.clearAllMocks()
vi.mocked(validateToolUse).mockImplementation(() => undefined)

mockTask = {
taskId: "test-task-id",
instanceId: "test-instance",
abort: false,
presentAssistantMessageLocked: false,
presentAssistantMessageHasPendingUpdates: false,
currentStreamingContentIndex: 0,
assistantMessageContent: [],
userMessageContent: [],
didCompleteReadingStream: false,
didRejectTool: false,
didAlreadyUseTool: false,
consecutiveMistakeCount: 0,
clineMessages: [],
api: {
getModel: () => ({ id: "test-model", info: {} }),
},
recordToolUsage: vi.fn(),
recordToolError: vi.fn(),
toolRepetitionDetector: {
check: vi.fn().mockReturnValue({ allowExecution: true }),
},
providerRef: {
deref: () => ({
getState: vi.fn().mockResolvedValue({
mode: "code",
customModes: [],
}),
}),
},
say: vi.fn().mockResolvedValue(undefined),
ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
pushToolResultToUserContent: vi.fn(),
}

mockTask.pushToolResultToUserContent = vi
.fn()
.mockImplementation((toolResult: Anthropic.ToolResultBlockParam) => {
const existingResult = mockTask.userMessageContent.find(
(block) => block.type === "tool_result" && block.tool_use_id === toolResult.tool_use_id,
)
if (existingResult) {
return false
}
mockTask.userMessageContent.push(toolResult)
return true
})
})

it("records exactly one attempt for a normal static tool", async () => {
mockTask.assistantMessageContent = [
{
type: "tool_use",
id: "call_1",
name: "read_file",
params: { path: "test.txt" },
nativeArgs: { path: "test.txt" },
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

expect(mockTask.recordToolUsage).toHaveBeenCalledTimes(1)
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("read_file")
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledTimes(1)
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "read_file")
})

it("records a valid dynamic mcp_ tool name as use_mcp_tool", async () => {
mockTask.assistantMessageContent = [
{
type: "tool_use",
id: "call_mcp",
name: "mcp_my_server_do_thing",
params: {},
nativeArgs: {},
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "use_mcp_tool")
})

it("records a malformed mcp_ tool name as use_mcp_tool, not the raw name", async () => {
mockTask.assistantMessageContent = [
{
type: "tool_use",
id: "call_mcp_bad",
name: "mcp_",
params: {},
nativeArgs: {},
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
expect(mockTask.recordToolUsage).not.toHaveBeenCalledWith("mcp_")
})

it("records a safe failure key without leaking the raw tool name when validation fails", async () => {
vi.mocked(validateToolUse).mockImplementation(() => {
throw new Error('Tool "read_file" is not allowed in this mode.')
})

mockTask.assistantMessageContent = [
{
type: "tool_use",
id: "call_bad_mode",
name: "read_file",
params: { path: "test.txt" },
nativeArgs: { path: "test.txt" },
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

// A known static tool that fails validation still maps to its own name
// (it's a real, recognized tool - just disallowed here), never left raw/unmapped.
expect(mockTask.recordToolError).toHaveBeenCalledWith("read_file", expect.any(String))
// No success attempt should be recorded for a validation failure.
expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
})

it("records invalid_tool_call, not the raw name, when an arbitrary unknown tool fails validation", async () => {
vi.mocked(validateToolUse).mockImplementation(() => {
throw new Error('Unknown tool "totally_made_up_tool". This tool does not exist.')
})

mockTask.assistantMessageContent = [
{
type: "tool_use",
id: "call_unknown",
name: "totally_made_up_tool",
params: {},
nativeArgs: {},
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

expect(mockTask.recordToolError).toHaveBeenCalledWith("invalid_tool_call", expect.any(String))
expect(mockTask.recordToolError).not.toHaveBeenCalledWith("totally_made_up_tool", expect.anything())
expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
})

describe("native mcp_tool_use block", () => {
it("records exactly one attempt once the MCP tool's own validation passes", async () => {
mockTask.providerRef = {
deref: () => ({
getState: vi.fn().mockResolvedValue({
mode: "code",
customModes: [],
}),
getMcpHub: () => ({
findServerNameBySanitizedName: () => "my_server",
}),
}),
}

mockTask.assistantMessageContent = [
{
type: "mcp_tool_use",
id: "call_native_mcp",
name: "mcp_my_server_do_thing",
serverName: "my_server",
toolName: "do_thing",
arguments: {},
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

expect(mockTask.recordToolUsage).toHaveBeenCalledTimes(1)
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledTimes(1)
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "use_mcp_tool")
})

it("records no attempt when the MCP server is not on the mode's allow-list", async () => {
vi.mocked(getModeBySlug).mockReturnValueOnce({
slug: "code",
name: "Code",
roleDefinition: "",
groups: [],
allowedMcpServers: ["some-other-server"],
})

mockTask.providerRef = {
deref: () => ({
getState: vi.fn().mockResolvedValue({
mode: "code",
customModes: [],
}),
getMcpHub: () => ({
findServerNameBySanitizedName: () => "my_server",
}),
}),
}

mockTask.assistantMessageContent = [
{
type: "mcp_tool_use",
id: "call_native_mcp_disallowed",
name: "mcp_my_server_do_thing",
serverName: "my_server",
toolName: "do_thing",
arguments: {},
partial: false,
},
]

await presentAssistantMessage(mockTask as unknown as Task)

// The server is disallowed, so the call never reaches onValidated:
// no success attempt is recorded for a call that was never permitted to execute.
expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
expect(TelemetryService.instance.captureToolUsage).not.toHaveBeenCalled()
})
})
})
Original file line number Diff line number Diff line change
Expand Up @@ -101,9 +101,10 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
// Verify consecutiveMistakeCount was incremented
expect(mockTask.consecutiveMistakeCount).toBe(1)

// Verify recordToolError was called
// Verify recordToolError was called with a safe static key, never the
// raw model-controlled tool name.
expect(mockTask.recordToolError).toHaveBeenCalledWith(
"nonexistent_tool",
"invalid_tool_call",
expect.stringContaining("Unknown tool"),
)

Expand Down Expand Up @@ -135,8 +136,9 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
// Verify consecutiveMistakeCount was incremented
expect(mockTask.consecutiveMistakeCount).toBe(1)

// Verify recordToolError was called
expect(mockTask.recordToolError).toHaveBeenCalled()
// Verify recordToolError was called with a safe static key, never the
// raw model-reported tool name ("fake_tool_that_does_not_exist").
expect(mockTask.recordToolError).toHaveBeenCalledWith("invalid_tool_call", expect.anything())

// Verify error message was shown to user
expect(mockTask.say).toHaveBeenCalledWith("error", expect.anything())
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
// npx vitest src/core/assistant-message/__tests__/toTelemetryToolName.spec.ts

import { describe, it, expect, vi } from "vitest"

vi.mock("../../tools/validateToolUse", () => ({
validateToolUse: vi.fn(),
isValidToolName: vi.fn((toolName: string) =>
["read_file", "write_to_file", "ask_followup_question", "attempt_completion", "use_mcp_tool"].includes(
toolName,
),
),
}))

import { toTelemetryToolName } from "../presentAssistantMessage"

describe("toTelemetryToolName", () => {
it("maps a known static tool to its own name", () => {
expect(toTelemetryToolName("read_file", false, undefined)).toBe("read_file")
})

it("maps a registered custom tool to custom_tool", () => {
expect(toTelemetryToolName("my_custom_tool", true, undefined)).toBe("custom_tool")
})

it("maps a valid dynamic mcp_ tool name to use_mcp_tool", () => {
expect(toTelemetryToolName("mcp_my_server_do_thing", false, undefined)).toBe("use_mcp_tool")
})

it("maps a malformed mcp_ tool name to use_mcp_tool", () => {
expect(toTelemetryToolName("mcp_", false, undefined)).toBe("use_mcp_tool")
})

it("maps an arbitrary unknown tool name to invalid_tool_call", () => {
expect(toTelemetryToolName("drop_table_users", false, undefined)).toBe("invalid_tool_call")
})

it("never returns the raw name for an unrecognized tool", () => {
const raw = "'; DROP TABLE users; --"
const result = toTelemetryToolName(raw, false, undefined)
expect(result).not.toBe(raw)
expect(result).toBe("invalid_tool_call")
})
})
Loading
Loading