Fix bug with role definition overrides for built-in modes

src/core/prompts/__tests__/system.test.ts

@@ -350,6 +350,56 @@ describe("SYSTEM_PROMPT", () => {
 		expect(customInstructionsIndex).toBeGreaterThan(userInstructionsHeader)
 	})
 
+	it("should use promptComponent roleDefinition when available", async () => {
+		const customPrompts = {
+			[defaultModeSlug]: {
+				roleDefinition: "Custom prompt role definition",
+				customInstructions: "Custom prompt instructions",
+			},
+		}
+
+		const prompt = await SYSTEM_PROMPT(
+			mockContext,
+			"/test/path",
+			false,
+			undefined,
+			undefined,
+			undefined,
+			defaultModeSlug,
+			customPrompts,
+			undefined,
+		)
+
+		// Role definition from promptComponent should be at the top
+		expect(prompt.indexOf("Custom prompt role definition")).toBeLessThan(prompt.indexOf("TOOL USE"))
+		// Should not contain the default mode's role definition
+		expect(prompt).not.toContain(modes[0].roleDefinition)
+	})
+
+	it("should fallback to modeConfig roleDefinition when promptComponent has no roleDefinition", async () => {
+		const customPrompts = {
+			[defaultModeSlug]: {
+				customInstructions: "Custom prompt instructions",
+				// No roleDefinition provided
+			},
+		}
+
+		const prompt = await SYSTEM_PROMPT(
+			mockContext,
+			"/test/path",
+			false,
+			undefined,
+			undefined,
+			undefined,
+			defaultModeSlug,
+			customPrompts,
+			undefined,
+		)
+
+		// Should use the default mode's role definition
+		expect(prompt.indexOf(modes[0].roleDefinition)).toBeLessThan(prompt.indexOf("TOOL USE"))
+	})
+
 	afterAll(() => {
 		jest.restoreAllMocks()
 	})

src/core/prompts/system.ts

@@ -54,7 +54,7 @@ async function generatePrompt(
 
 	// Get the full mode config to ensure we have the role definition
 	const modeConfig = getModeBySlug(mode, customModeConfigs) || modes.find((m) => m.slug === mode) || modes[0]
-	const roleDefinition = modeConfig.roleDefinition
+	const roleDefinition = promptComponent?.roleDefinition || modeConfig.roleDefinition
 
 	const basePrompt = `${roleDefinition}