Skip to content

Commit e8feed7

Browse files
ChuKhaLiChuKhaLi
committed
Refactors mode selection logic for custom modes
Refactors the mode selection logic to prioritize custom modes
1 parent 2698f78 commit e8feed7

File tree

3 files changed

+130
-75
lines changed

3 files changed

+130
-75
lines changed

‎src/shared/__tests__/modes.test.ts‎

Lines changed: 92 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// npx jest src/shared/__tests__/modes.test.ts
22

3-
import type { ModeConfig } from "@roo-code/types"
3+
import type { ModeConfig, PromptComponent } from "@roo-code/types"
44

55
// Mock setup must come before imports
66
jest.mock("vscode")
@@ -374,7 +374,6 @@ describe("FileRestrictionError", () => {
374374

375375
describe("getModeSelection", () => {
376376
const builtInAskMode = modes.find((m) => m.slug === "ask")!
377-
378377
const customModesList: ModeConfig[] = [
379378
{
380379
slug: "code", // Override
@@ -408,8 +407,8 @@ describe("getModeSelection", () => {
408407
expect(selection.baseInstructions).toBe(builtInAskMode.customInstructions || "")
409408
})
410409

411-
test("should prioritize promptComponent for built-in mode", () => {
412-
const selection = getModeSelection("ask", promptComponentAsk)
410+
test("should prioritize promptComponent for built-in mode if no custom mode exists for that slug", () => {
411+
const selection = getModeSelection("ask", promptComponentAsk) // "ask" is not in customModesList
413412
expect(selection.roleDefinition).toBe(promptComponentAsk.roleDefinition)
414413
expect(selection.baseInstructions).toBe(promptComponentAsk.customInstructions)
415414
})
@@ -421,21 +420,21 @@ describe("getModeSelection", () => {
421420
expect(selection.baseInstructions).toBe(customCode.customInstructions)
422421
})
423422

424-
test("should prioritize customMode over promptComponent and built-in mode when all define properties", () => {
423+
test("should prioritize customMode over promptComponent and built-in mode", () => {
425424
const selection = getModeSelection("code", promptComponentCode, customModesList)
426425
const customCode = customModesList.find((m) => m.slug === "code")!
427426
expect(selection.roleDefinition).toBe(customCode.roleDefinition)
428427
expect(selection.baseInstructions).toBe(customCode.customInstructions)
429428
})
430429

431-
test("should return new custom mode details", () => {
430+
test("should return new custom mode details if it exists", () => {
432431
const selection = getModeSelection("new-custom", undefined, customModesList)
433432
const newCustom = customModesList.find((m) => m.slug === "new-custom")!
434433
expect(selection.roleDefinition).toBe(newCustom.roleDefinition)
435434
expect(selection.baseInstructions).toBe(newCustom.customInstructions)
436435
})
437436

438-
test("customMode properties take precedence for new custom mode even with promptComponent", () => {
437+
test("customMode takes precedence for a new custom mode even if promptComponent is provided", () => {
439438
const promptComponentNew: PromptComponent = {
440439
roleDefinition: "Prompt New Custom Role",
441440
customInstructions: "Prompt New Custom Instructions",
@@ -446,33 +445,24 @@ describe("getModeSelection", () => {
446445
expect(selection.baseInstructions).toBe(newCustomMode.customInstructions)
447446
})
448447

449-
test("should fallback to default (first) mode if slug does not exist", () => {
448+
test("should return empty strings if slug does not exist in custom, prompt, or built-in modes", () => {
450449
const selection = getModeSelection("non-existent-mode", undefined, customModesList)
451-
expect(selection.roleDefinition).toBe(modes[0].roleDefinition)
452-
expect(selection.baseInstructions).toBe(modes[0].customInstructions || "")
453-
})
454-
455-
test("customMode.roleDefinition is used if defined, ignoring promptComponent's; customMode.customInstructions also from customMode", () => {
456-
const selection = getModeSelection("code", { roleDefinition: "Prompt Role Only" }, customModesList)
457-
const customCodeMode = customModesList.find((m) => m.slug === "code")!
458-
expect(selection.roleDefinition).toBe(customCodeMode.roleDefinition)
459-
expect(selection.baseInstructions).toBe(customCodeMode.customInstructions)
450+
expect(selection.roleDefinition).toBe("")
451+
expect(selection.baseInstructions).toBe("")
460452
})
461453

462-
test("customMode.customInstructions is used if defined, ignoring promptComponent's; customMode.roleDefinition also from customMode", () => {
463-
const selection = getModeSelection("code", { customInstructions: "Prompt Instructions Only" }, customModesList)
454+
test("customMode's properties are used if customMode exists, ignoring promptComponent's properties", () => {
455+
const selection = getModeSelection(
456+
"code",
457+
{ roleDefinition: "Prompt Role Only", customInstructions: "Prompt Instructions Only" },
458+
customModesList,
459+
)
464460
const customCodeMode = customModesList.find((m) => m.slug === "code")!
465-
expect(selection.roleDefinition).toBe(customCodeMode.roleDefinition)
466-
expect(selection.baseInstructions).toBe(customCodeMode.customInstructions)
467-
})
468-
469-
test("customMode takes precedence over built-in when no promptComponent", () => {
470-
const selection = getModeSelection("code", undefined, customModesList)
471-
expect(selection.roleDefinition).toBe(customModesList.find((m) => m.slug === "code")!.roleDefinition)
472-
expect(selection.baseInstructions).toBe(customModesList.find((m) => m.slug === "code")!.customInstructions)
461+
expect(selection.roleDefinition).toBe(customCodeMode.roleDefinition) // Takes from customCodeMode
462+
expect(selection.baseInstructions).toBe(customCodeMode.customInstructions) // Takes from customCodeMode
473463
})
474464

475-
test("handles undefined customInstructions in modes gracefully", () => {
465+
test("handles undefined customInstructions in customMode gracefully", () => {
476466
const modesWithoutCustomInstructions: ModeConfig[] = [
477467
{
478468
slug: "no-instr",
@@ -484,65 +474,106 @@ describe("getModeSelection", () => {
484474
]
485475
const selection = getModeSelection("no-instr", undefined, modesWithoutCustomInstructions)
486476
expect(selection.roleDefinition).toBe("Role for no instructions")
487-
expect(selection.baseInstructions).toBe("") // Should default to empty string
477+
expect(selection.baseInstructions).toBe("") // Defaults to empty string
488478
})
489479

490-
test("handles undefined roleDefinition in modes gracefully by falling back", () => {
491-
const modesWithoutRoleDef: ModeConfig[] = [
480+
test("handles empty or undefined roleDefinition in customMode gracefully", () => {
481+
const modesWithEmptyRoleDef: ModeConfig[] = [
492482
{
493-
slug: "no-role",
494-
name: "No Role Mode",
495-
roleDefinition: "", // Ensure roleDefinition is present
496-
customInstructions: "Instructions for no role",
483+
slug: "empty-role",
484+
name: "Empty Role Mode",
485+
roleDefinition: "",
486+
customInstructions: "Instructions for empty role",
497487
groups: ["read"],
498488
},
499489
]
500-
// Since 'no-role' is a custom mode not overriding a built-in, and it lacks a role,
501-
// the logic in getModeSelection might fall back if not handled.
502-
// The current getModeBySlug logic in getModeSelection will fetch this mode.
503-
// Then, isCustom will be true.
504-
// roleDefinition = modeConfig?.roleDefinition (undefined) || promptComponent?.roleDefinition (undefined) || ""
505-
// So it should become ""
506-
const selection = getModeSelection("no-role", undefined, modesWithoutRoleDef)
490+
const selection = getModeSelection("empty-role", undefined, modesWithEmptyRoleDef)
507491
expect(selection.roleDefinition).toBe("")
508-
expect(selection.baseInstructions).toBe("Instructions for no role")
492+
expect(selection.baseInstructions).toBe("Instructions for empty role")
493+
494+
const modesWithUndefinedRoleDef: ModeConfig[] = [
495+
{
496+
slug: "undefined-role",
497+
name: "Undefined Role Mode",
498+
roleDefinition: "", // Test undefined explicitly by using an empty string
499+
customInstructions: "Instructions for undefined role",
500+
groups: ["read"],
501+
},
502+
]
503+
const selection2 = getModeSelection("undefined-role", undefined, modesWithUndefinedRoleDef)
504+
expect(selection2.roleDefinition).toBe("")
505+
expect(selection2.baseInstructions).toBe("Instructions for undefined role")
509506
})
510507

511-
test("promptComponent fills customInstructions if customMode's is undefined", () => {
512-
const customModeRoleOnly: ModeConfig[] = [
513-
{ slug: "role-custom", name: "Role Custom", roleDefinition: "Custom Role Only", groups: ["read"] },
508+
test("customMode's defined properties take precedence, undefined ones in customMode result in ''", () => {
509+
const customModeRoleOnlyList: ModeConfig[] = [
510+
// Renamed for clarity
511+
{
512+
slug: "role-custom",
513+
name: "Role Custom",
514+
roleDefinition: "Custom Role Only",
515+
groups: ["read"] /* customInstructions undefined */,
516+
},
514517
]
515518
const promptComponentInstrOnly: PromptComponent = { customInstructions: "Prompt Instructions Only" }
516-
const selection = getModeSelection("role-custom", promptComponentInstrOnly, customModeRoleOnly)
517-
expect(selection.roleDefinition).toBe("Custom Role Only")
518-
expect(selection.baseInstructions).toBe("Prompt Instructions Only")
519+
// "role-custom" exists in customModeRoleOnlyList
520+
const selection = getModeSelection("role-custom", promptComponentInstrOnly, customModeRoleOnlyList)
521+
// customMode is chosen.
522+
expect(selection.roleDefinition).toBe("Custom Role Only") // From customMode
523+
expect(selection.baseInstructions).toBe("") // From customMode (undefined || '' -> '')
519524
})
520525

521-
test("promptComponent fills roleDefinition if customMode's is undefined", () => {
522-
const customModeInstrOnly: ModeConfig[] = [
526+
test("customMode's defined properties take precedence, empty string ones in customMode are used", () => {
527+
const customModeInstrOnlyList: ModeConfig[] = [
528+
// Renamed for clarity
523529
{
524530
slug: "instr-custom",
525531
name: "Instr Custom",
526-
roleDefinition: "", // Added to satisfy ModeConfig type
532+
roleDefinition: "", // Explicitly empty
527533
customInstructions: "Custom Instructions Only",
528534
groups: ["read"],
529535
},
530536
]
531537
const promptComponentRoleOnly: PromptComponent = { roleDefinition: "Prompt Role Only" }
532-
const selection = getModeSelection("instr-custom", promptComponentRoleOnly, customModeInstrOnly)
533-
expect(selection.roleDefinition).toBe("Prompt Role Only")
534-
expect(selection.baseInstructions).toBe("Custom Instructions Only")
538+
// "instr-custom" exists in customModeInstrOnlyList
539+
const selection = getModeSelection("instr-custom", promptComponentRoleOnly, customModeInstrOnlyList)
540+
// customMode is chosen
541+
expect(selection.roleDefinition).toBe("") // From customMode ( "" || '' -> "")
542+
expect(selection.baseInstructions).toBe("Custom Instructions Only") // From customMode
535543
})
536544

537-
test("builtInMode fills properties if customMode and promptComponent's are undefined", () => {
545+
test("customMode with empty/undefined fields takes precedence over promptComponent and builtInMode", () => {
538546
const customModeMinimal: ModeConfig[] = [
539-
{ slug: "ask", name: "Custom Ask Minimal", roleDefinition: "", groups: ["read"] }, // No roleDef or customInstr
547+
{ slug: "ask", name: "Custom Ask Minimal", roleDefinition: "", groups: ["read"] }, // roleDef empty, customInstr undefined
540548
]
541-
const promptComponentMinimal: PromptComponent = {} // No roleDef or customInstr
549+
const promptComponentMinimal: PromptComponent = {
550+
roleDefinition: "Prompt Min Role",
551+
customInstructions: "Prompt Min Instr",
552+
}
553+
// "ask" is in customModeMinimal
542554
const selection = getModeSelection("ask", promptComponentMinimal, customModeMinimal)
555+
// customMode is chosen
556+
expect(selection.roleDefinition).toBe("") // From customModeMinimal
557+
expect(selection.baseInstructions).toBe("") // From customModeMinimal
558+
})
559+
560+
test("promptComponent is used if customMode for slug does not exist, even if customModesList is provided", () => {
561+
// 'ask' is not in customModesList, but 'code' and 'new-custom' are.
562+
const selection = getModeSelection("ask", promptComponentAsk, customModesList)
563+
expect(selection.roleDefinition).toBe(promptComponentAsk.roleDefinition)
564+
expect(selection.baseInstructions).toBe(promptComponentAsk.customInstructions)
565+
})
543566

544-
// According to original logic, if customMode provides "", that takes precedence.
545-
expect(selection.roleDefinition).toBe("") // Was builtInAskMode.roleDefinition
546-
expect(selection.baseInstructions).toBe("") // Was builtInAskMode.customInstructions || ""
567+
test("builtInMode is used if customMode for slug does not exist and promptComponent is not provided", () => {
568+
// 'ask' is not in customModesList
569+
const selection = getModeSelection("ask", undefined, customModesList)
570+
expect(selection.roleDefinition).toBe(builtInAskMode.roleDefinition)
571+
expect(selection.baseInstructions).toBe(builtInAskMode.customInstructions || "")
572+
})
573+
574+
test("promptComponent is used if customMode is not provided (undefined customModesList)", () => {
575+
const selection = getModeSelection("ask", promptComponentAsk, undefined)
576+
expect(selection.roleDefinition).toBe(promptComponentAsk.roleDefinition)
577+
expect(selection.baseInstructions).toBe(promptComponentAsk.customInstructions)
547578
})
548579
})

‎src/shared/modes.ts‎

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
import * as vscode from "vscode"
22

3-
import type { GroupOptions, GroupEntry, ModeConfig, CustomModePrompts, ExperimentId, ToolGroup } from "@roo-code/types"
3+
import type {
4+
GroupOptions,
5+
GroupEntry,
6+
ModeConfig,
7+
CustomModePrompts,
8+
ExperimentId,
9+
ToolGroup,
10+
PromptComponent,
11+
} from "@roo-code/types"
412

513
import { addCustomInstructions } from "../core/prompts/sections/custom-instructions"
614

@@ -149,16 +157,27 @@ export function isCustomMode(slug: string, customModes?: ModeConfig[]): boolean
149157
return !!customModes?.some((mode) => mode.slug === slug)
150158
}
151159

160+
/**
161+
* Find a mode by its slug, don't fall back to built-in modes
162+
*/
163+
export function findModeBySlug(slug: string, modes: readonly ModeConfig[] | undefined): ModeConfig | undefined {
164+
return modes?.find((mode) => mode.slug === slug)
165+
}
166+
167+
/**
168+
* Get the mode selection based on the provided mode slug, prompt component, and custom modes.
169+
* If a custom mode is found, it takes precedence over the built-in modes.
170+
* If no custom mode is found, the built-in mode is used.
171+
* If neither is found, the default mode is used.
172+
*/
152173
export function getModeSelection(mode: string, promptComponent?: PromptComponent, customModes?: ModeConfig[]) {
153-
const modeConfig = getModeBySlug(mode, customModes) || modes.find((m) => m.slug === mode) || modes[0]
154-
const isCustom = isCustomMode(mode, customModes)
155-
const roleDefinition = isCustom
156-
? modeConfig?.roleDefinition || promptComponent?.roleDefinition || ""
157-
: promptComponent?.roleDefinition || modeConfig.roleDefinition || ""
158-
159-
const baseInstructions = isCustom
160-
? modeConfig?.customInstructions || promptComponent?.customInstructions || ""
161-
: promptComponent?.customInstructions || modeConfig.customInstructions || ""
174+
const customMode = findModeBySlug(mode, customModes)
175+
const builtInMode = findModeBySlug(mode, modes)
176+
177+
const modeToUse = customMode || promptComponent || builtInMode
178+
179+
const roleDefinition = modeToUse?.roleDefinition || ""
180+
const baseInstructions = modeToUse?.customInstructions || ""
162181

163182
return {
164183
roleDefinition,

‎webview-ui/src/components/prompts/PromptsView.tsx‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,14 @@ import { ChevronsUpDown, X } from "lucide-react"
1111

1212
import { ModeConfig, GroupEntry, PromptComponent, ToolGroup, modeConfigSchema } from "@roo-code/types"
1313

14-
import { Mode, getRoleDefinition, getWhenToUse, getCustomInstructions, getAllModes } from "@roo/modes"
14+
import {
15+
Mode,
16+
getRoleDefinition,
17+
getWhenToUse,
18+
getCustomInstructions,
19+
getAllModes,
20+
findModeBySlug as findCustomModeBySlug,
21+
} from "@roo/modes"
1522
import { supportPrompt, SupportPromptType } from "@roo/support-prompt"
1623
import { TOOL_GROUPS } from "@roo/tools"
1724

@@ -133,9 +140,7 @@ const PromptsView = ({ onDone }: PromptsViewProps) => {
133140
// Helper function to find a mode by slug
134141
const findModeBySlug = useCallback(
135142
(searchSlug: string, modes: readonly ModeConfig[] | undefined): ModeConfig | undefined => {
136-
if (!modes) return undefined
137-
const isModeWithSlug = (mode: ModeConfig): mode is ModeConfig => mode.slug === searchSlug
138-
return modes.find(isModeWithSlug)
143+
return findCustomModeBySlug(searchSlug, modes)
139144
},
140145
[],
141146
)

0 commit comments

Comments
 (0)