fix: ensure correct precedence for roleDefinition and customInstructions

ChuKhaLi · ChuKhaLi · commit 2698f78c2017 · 2025-05-28T06:45:18.000+07:00
diff --git a/src/core/prompts/system.ts b/src/core/prompts/system.ts
@@ -3,7 +3,7 @@ import * as os from "os"
 
 import type { ModeConfig, PromptComponent, CustomModePrompts } from "@roo-code/types"
 
-import { Mode, modes, defaultModeSlug, getModeBySlug, getGroupName } from "../../shared/modes"
+import { Mode, modes, defaultModeSlug, getModeBySlug, getGroupName, getModeSelection } from "../../shared/modes"
 import { DiffStrategy } from "../../shared/tools"
 import { formatLanguage } from "../../shared/language"
 
@@ -51,9 +51,9 @@ async function generatePrompt(
 	// If diff is disabled, don't pass the diffStrategy
 	const effectiveDiffStrategy = diffEnabled ? diffStrategy : undefined
 
-	// Get the full mode config to ensure we have the role definition
+	// Get the full mode config to ensure we have the role definition (used for groups, etc.)
 	const modeConfig = getModeBySlug(mode, customModeConfigs) || modes.find((m) => m.slug === mode) || modes[0]
-	const roleDefinition = promptComponent?.roleDefinition || modeConfig.roleDefinition
+	const { roleDefinition, baseInstructions } = getModeSelection(mode, promptComponent, customModeConfigs)
 
 	const [modesSection, mcpServersSection] = await Promise.all([
 		getModesSection(context),
@@ -97,7 +97,7 @@ ${getSystemInfoSection(cwd)}
 
 ${getObjectiveSection()}
 
-${await addCustomInstructions(promptComponent?.customInstructions || modeConfig.customInstructions || "", globalCustomInstructions || "", cwd, mode, { language: language ?? formatLanguage(vscode.env.language), rooIgnoreInstructions })}`
+${await addCustomInstructions(baseInstructions, globalCustomInstructions || "", cwd, mode, { language: language ?? formatLanguage(vscode.env.language), rooIgnoreInstructions })}`
 
 	return basePrompt
 }
@@ -149,9 +149,14 @@ export const SYSTEM_PROMPT = async (
 
 	// If a file-based custom system prompt exists, use it
 	if (fileCustomSystemPrompt) {
-		const roleDefinition = promptComponent?.roleDefinition || currentMode.roleDefinition
+		const { roleDefinition, baseInstructions: baseInstructionsForFile } = getModeSelection(
+			mode,
+			promptComponent,
+			customModes,
+		)
+
 		const customInstructions = await addCustomInstructions(
-			promptComponent?.customInstructions || currentMode.customInstructions || "",
+			baseInstructionsForFile,
 			globalCustomInstructions || "",
 			cwd,
 			mode,
diff --git a/src/shared/__tests__/modes.test.ts b/src/shared/__tests__/modes.test.ts
@@ -11,7 +11,7 @@ jest.mock("../../core/prompts/sections/custom-instructions", () => ({
 	addCustomInstructions: mockAddCustomInstructions,
 }))
 
-import { isToolAllowedForMode, FileRestrictionError, getFullModeDetails, modes } from "../modes"
+import { isToolAllowedForMode, FileRestrictionError, getFullModeDetails, modes, getModeSelection } from "../modes"
 import { addCustomInstructions } from "../../core/prompts/sections/custom-instructions"
 
 describe("isToolAllowedForMode", () => {
@@ -371,3 +371,178 @@ describe("FileRestrictionError", () => {
 		expect(error.name).toBe("FileRestrictionError")
 	})
 })
+
+describe("getModeSelection", () => {
+	const builtInAskMode = modes.find((m) => m.slug === "ask")!
+
+	const customModesList: ModeConfig[] = [
+		{
+			slug: "code", // Override
+			name: "Custom Code Mode",
+			roleDefinition: "Custom Code Role",
+			customInstructions: "Custom Code Instructions",
+			groups: ["read"],
+		},
+		{
+			slug: "new-custom",
+			name: "New Custom Mode",
+			roleDefinition: "New Custom Role",
+			customInstructions: "New Custom Instructions",
+			groups: ["edit"],
+		},
+	]
+
+	const promptComponentCode: PromptComponent = {
+		roleDefinition: "Prompt Component Code Role",
+		customInstructions: "Prompt Component Code Instructions",
+	}
+
+	const promptComponentAsk: PromptComponent = {
+		roleDefinition: "Prompt Component Ask Role",
+		customInstructions: "Prompt Component Ask Instructions",
+	}
+
+	test("should return built-in mode details if no overrides", () => {
+		const selection = getModeSelection("ask")
+		expect(selection.roleDefinition).toBe(builtInAskMode.roleDefinition)
+		expect(selection.baseInstructions).toBe(builtInAskMode.customInstructions || "")
+	})
+
+	test("should prioritize promptComponent for built-in mode", () => {
+		const selection = getModeSelection("ask", promptComponentAsk)
+		expect(selection.roleDefinition).toBe(promptComponentAsk.roleDefinition)
+		expect(selection.baseInstructions).toBe(promptComponentAsk.customInstructions)
+	})
+
+	test("should prioritize customMode over built-in mode", () => {
+		const selection = getModeSelection("code", undefined, customModesList)
+		const customCode = customModesList.find((m) => m.slug === "code")!
+		expect(selection.roleDefinition).toBe(customCode.roleDefinition)
+		expect(selection.baseInstructions).toBe(customCode.customInstructions)
+	})
+
+	test("should prioritize customMode over promptComponent and built-in mode when all define properties", () => {
+		const selection = getModeSelection("code", promptComponentCode, customModesList)
+		const customCode = customModesList.find((m) => m.slug === "code")!
+		expect(selection.roleDefinition).toBe(customCode.roleDefinition)
+		expect(selection.baseInstructions).toBe(customCode.customInstructions)
+	})
+
+	test("should return new custom mode details", () => {
+		const selection = getModeSelection("new-custom", undefined, customModesList)
+		const newCustom = customModesList.find((m) => m.slug === "new-custom")!
+		expect(selection.roleDefinition).toBe(newCustom.roleDefinition)
+		expect(selection.baseInstructions).toBe(newCustom.customInstructions)
+	})
+
+	test("customMode properties take precedence for new custom mode even with promptComponent", () => {
+		const promptComponentNew: PromptComponent = {
+			roleDefinition: "Prompt New Custom Role",
+			customInstructions: "Prompt New Custom Instructions",
+		}
+		const selection = getModeSelection("new-custom", promptComponentNew, customModesList)
+		const newCustomMode = customModesList.find((m) => m.slug === "new-custom")!
+		expect(selection.roleDefinition).toBe(newCustomMode.roleDefinition)
+		expect(selection.baseInstructions).toBe(newCustomMode.customInstructions)
+	})
+
+	test("should fallback to default (first) mode if slug does not exist", () => {
+		const selection = getModeSelection("non-existent-mode", undefined, customModesList)
+		expect(selection.roleDefinition).toBe(modes[0].roleDefinition)
+		expect(selection.baseInstructions).toBe(modes[0].customInstructions || "")
+	})
+
+	test("customMode.roleDefinition is used if defined, ignoring promptComponent's; customMode.customInstructions also from customMode", () => {
+		const selection = getModeSelection("code", { roleDefinition: "Prompt Role Only" }, customModesList)
+		const customCodeMode = customModesList.find((m) => m.slug === "code")!
+		expect(selection.roleDefinition).toBe(customCodeMode.roleDefinition)
+		expect(selection.baseInstructions).toBe(customCodeMode.customInstructions)
+	})
+
+	test("customMode.customInstructions is used if defined, ignoring promptComponent's; customMode.roleDefinition also from customMode", () => {
+		const selection = getModeSelection("code", { customInstructions: "Prompt Instructions Only" }, customModesList)
+		const customCodeMode = customModesList.find((m) => m.slug === "code")!
+		expect(selection.roleDefinition).toBe(customCodeMode.roleDefinition)
+		expect(selection.baseInstructions).toBe(customCodeMode.customInstructions)
+	})
+
+	test("customMode takes precedence over built-in when no promptComponent", () => {
+		const selection = getModeSelection("code", undefined, customModesList)
+		expect(selection.roleDefinition).toBe(customModesList.find((m) => m.slug === "code")!.roleDefinition)
+		expect(selection.baseInstructions).toBe(customModesList.find((m) => m.slug === "code")!.customInstructions)
+	})
+
+	test("handles undefined customInstructions in modes gracefully", () => {
+		const modesWithoutCustomInstructions: ModeConfig[] = [
+			{
+				slug: "no-instr",
+				name: "No Instructions Mode",
+				roleDefinition: "Role for no instructions",
+				groups: ["read"],
+				// customInstructions is undefined
+			},
+		]
+		const selection = getModeSelection("no-instr", undefined, modesWithoutCustomInstructions)
+		expect(selection.roleDefinition).toBe("Role for no instructions")
+		expect(selection.baseInstructions).toBe("") // Should default to empty string
+	})
+
+	test("handles undefined roleDefinition in modes gracefully by falling back", () => {
+		const modesWithoutRoleDef: ModeConfig[] = [
+			{
+				slug: "no-role",
+				name: "No Role Mode",
+				roleDefinition: "", // Ensure roleDefinition is present
+				customInstructions: "Instructions for no role",
+				groups: ["read"],
+			},
+		]
+		// Since 'no-role' is a custom mode not overriding a built-in, and it lacks a role,
+		// the logic in getModeSelection might fall back if not handled.
+		// The current getModeBySlug logic in getModeSelection will fetch this mode.
+		// Then, isCustom will be true.
+		// roleDefinition = modeConfig?.roleDefinition (undefined) || promptComponent?.roleDefinition (undefined) || ""
+		// So it should become ""
+		const selection = getModeSelection("no-role", undefined, modesWithoutRoleDef)
+		expect(selection.roleDefinition).toBe("")
+		expect(selection.baseInstructions).toBe("Instructions for no role")
+	})
+
+	test("promptComponent fills customInstructions if customMode's is undefined", () => {
+		const customModeRoleOnly: ModeConfig[] = [
+			{ slug: "role-custom", name: "Role Custom", roleDefinition: "Custom Role Only", groups: ["read"] },
+		]
+		const promptComponentInstrOnly: PromptComponent = { customInstructions: "Prompt Instructions Only" }
+		const selection = getModeSelection("role-custom", promptComponentInstrOnly, customModeRoleOnly)
+		expect(selection.roleDefinition).toBe("Custom Role Only")
+		expect(selection.baseInstructions).toBe("Prompt Instructions Only")
+	})
+
+	test("promptComponent fills roleDefinition if customMode's is undefined", () => {
+		const customModeInstrOnly: ModeConfig[] = [
+			{
+				slug: "instr-custom",
+				name: "Instr Custom",
+				roleDefinition: "", // Added to satisfy ModeConfig type
+				customInstructions: "Custom Instructions Only",
+				groups: ["read"],
+			},
+		]
+		const promptComponentRoleOnly: PromptComponent = { roleDefinition: "Prompt Role Only" }
+		const selection = getModeSelection("instr-custom", promptComponentRoleOnly, customModeInstrOnly)
+		expect(selection.roleDefinition).toBe("Prompt Role Only")
+		expect(selection.baseInstructions).toBe("Custom Instructions Only")
+	})
+
+	test("builtInMode fills properties if customMode and promptComponent's are undefined", () => {
+		const customModeMinimal: ModeConfig[] = [
+			{ slug: "ask", name: "Custom Ask Minimal", roleDefinition: "", groups: ["read"] }, // No roleDef or customInstr
+		]
+		const promptComponentMinimal: PromptComponent = {} // No roleDef or customInstr
+		const selection = getModeSelection("ask", promptComponentMinimal, customModeMinimal)
+
+		// According to original logic, if customMode provides "", that takes precedence.
+		expect(selection.roleDefinition).toBe("") // Was builtInAskMode.roleDefinition
+		expect(selection.baseInstructions).toBe("") // Was builtInAskMode.customInstructions || ""
+	})
+})
diff --git a/src/shared/modes.ts b/src/shared/modes.ts
@@ -149,6 +149,23 @@ export function isCustomMode(slug: string, customModes?: ModeConfig[]): boolean
 	return !!customModes?.some((mode) => mode.slug === slug)
 }
 
+export function getModeSelection(mode: string, promptComponent?: PromptComponent, customModes?: ModeConfig[]) {
+	const modeConfig = getModeBySlug(mode, customModes) || modes.find((m) => m.slug === mode) || modes[0]
+	const isCustom = isCustomMode(mode, customModes)
+	const roleDefinition = isCustom
+		? modeConfig?.roleDefinition || promptComponent?.roleDefinition || ""
+		: promptComponent?.roleDefinition || modeConfig.roleDefinition || ""
+
+	const baseInstructions = isCustom
+		? modeConfig?.customInstructions || promptComponent?.customInstructions || ""
+		: promptComponent?.customInstructions || modeConfig.customInstructions || ""
+
+	return {
+		roleDefinition,
+		baseInstructions,
+	}
+}
+
 // Custom error class for file restrictions
 export class FileRestrictionError extends Error {
 	constructor(mode: string, pattern: string, description: string | undefined, filePath: string) {
