Skip to content

Commit 74f9150

Browse files
roomoteroomote
committed
feat: allow users to change mode when approving orchestrator subtasks
- Added ModeSelector component to newTask tool approval UI in ChatRow - Users can now click on the mode to select a different one before approving - Selected mode is passed through the askResponse flow to the backend - Modified Task.ts to store and provide access to askResponseValues - Updated newTaskTool.ts to use the user-selected mode if provided - Added comprehensive tests for the new functionality - Added translation key for mode selector tooltip Fixes #6706
1 parent d90bab7 commit 74f9150

File tree

6 files changed

+239
-15
lines changed

6 files changed

+239
-15
lines changed

‎src/core/task/Task.ts‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
225225
private askResponse?: ClineAskResponse
226226
private askResponseText?: string
227227
private askResponseImages?: string[]
228+
private askResponseValues?: Record<string, any>
228229
public lastMessageTs?: number
229230

231+
// Getter for askResponseValues to allow tools to access it
232+
get getAskResponseValues(): Record<string, any> | undefined {
233+
return this.askResponseValues
234+
}
235+
230236
// Tool Use
231237
consecutiveMistakeCount: number = 0
232238
consecutiveMistakeLimit: number
@@ -742,10 +748,19 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
742748
this.handleWebviewAskResponse("messageResponse", text, images)
743749
}
744750

745-
handleWebviewAskResponse(askResponse: ClineAskResponse, text?: string, images?: string[]) {
751+
handleWebviewAskResponse(
752+
askResponse: ClineAskResponse,
753+
text?: string,
754+
images?: string[],
755+
values?: Record<string, any>,
756+
) {
746757
this.askResponse = askResponse
747758
this.askResponseText = text
748759
this.askResponseImages = images
760+
// Store values for later use if needed
761+
if (values) {
762+
this.askResponseValues = values
763+
}
749764
}
750765

751766
async handleTerminalOperation(terminalOperation: "continue" | "abort") {

‎src/core/tools/__tests__/newTaskTool.spec.ts‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ const mockCline = {
3636
consecutiveMistakeCount: 0,
3737
isPaused: false,
3838
pausedModeSlug: "ask",
39+
getAskResponseValues: undefined as Record<string, any> | undefined,
3940
providerRef: {
4041
deref: vi.fn(() => ({
4142
getState: vi.fn(() => ({ customModes: [], mode: "ask" })),
@@ -184,4 +185,143 @@ describe("newTaskTool", () => {
184185
})
185186

186187
// Add more tests for error handling (missing params, invalid mode, approval denied) if needed
188+
189+
it("should use user-selected mode when provided in askResponseValues", async () => {
190+
const block: ToolUse = {
191+
type: "tool_use",
192+
name: "new_task",
193+
params: {
194+
mode: "code",
195+
message: "Create a new feature",
196+
},
197+
partial: false,
198+
}
199+
200+
// Mock user selecting a different mode
201+
mockCline.getAskResponseValues = { selectedMode: "architect" }
202+
203+
// Mock the architect mode
204+
vi.mocked(getModeBySlug).mockImplementation((slug) => {
205+
if (slug === "architect") {
206+
return {
207+
slug: "architect",
208+
name: "Architect Mode",
209+
roleDefinition: "Architecture role definition",
210+
groups: ["command", "read"],
211+
}
212+
}
213+
return {
214+
slug: "code",
215+
name: "Code Mode",
216+
roleDefinition: "Test role definition",
217+
groups: ["command", "read", "edit"],
218+
}
219+
})
220+
221+
const mockHandleModeSwitch = vi.fn()
222+
mockCline.providerRef.deref = vi.fn(() => ({
223+
getState: vi.fn(() => ({ customModes: [], mode: "ask" })),
224+
handleModeSwitch: mockHandleModeSwitch,
225+
initClineWithTask: mockInitClineWithTask,
226+
}))
227+
228+
await newTaskTool(
229+
mockCline as any,
230+
block,
231+
mockAskApproval,
232+
mockHandleError,
233+
mockPushToolResult,
234+
mockRemoveClosingTag,
235+
)
236+
237+
// Verify the mode switch was called with the user-selected mode
238+
expect(mockHandleModeSwitch).toHaveBeenCalledWith("architect")
239+
240+
// Verify the success message includes the correct mode name
241+
expect(mockPushToolResult).toHaveBeenCalledWith(
242+
expect.stringContaining("Successfully created new task in Architect Mode"),
243+
)
244+
})
245+
246+
it("should use original mode when no user selection is provided", async () => {
247+
const block: ToolUse = {
248+
type: "tool_use",
249+
name: "new_task",
250+
params: {
251+
mode: "code",
252+
message: "Create a new feature",
253+
},
254+
partial: false,
255+
}
256+
257+
// No user selection
258+
mockCline.getAskResponseValues = undefined
259+
260+
const mockHandleModeSwitch = vi.fn()
261+
mockCline.providerRef.deref = vi.fn(() => ({
262+
getState: vi.fn(() => ({ customModes: [], mode: "ask" })),
263+
handleModeSwitch: mockHandleModeSwitch,
264+
initClineWithTask: mockInitClineWithTask,
265+
}))
266+
267+
await newTaskTool(
268+
mockCline as any,
269+
block,
270+
mockAskApproval,
271+
mockHandleError,
272+
mockPushToolResult,
273+
mockRemoveClosingTag,
274+
)
275+
276+
// Verify the mode switch was called with the original mode
277+
expect(mockHandleModeSwitch).toHaveBeenCalledWith("code")
278+
279+
// Verify the success message includes the correct mode name
280+
expect(mockPushToolResult).toHaveBeenCalledWith(
281+
expect.stringContaining("Successfully created new task in Code Mode"),
282+
)
283+
})
284+
285+
it("should handle invalid user-selected mode gracefully", async () => {
286+
const block: ToolUse = {
287+
type: "tool_use",
288+
name: "new_task",
289+
params: {
290+
mode: "code",
291+
message: "Create a new feature",
292+
},
293+
partial: false,
294+
}
295+
296+
// Mock user selecting an invalid mode
297+
mockCline.getAskResponseValues = { selectedMode: "invalid-mode" }
298+
299+
// Mock getModeBySlug to return undefined for invalid mode
300+
vi.mocked(getModeBySlug).mockImplementation((slug) => {
301+
if (slug === "invalid-mode") {
302+
return undefined
303+
}
304+
return {
305+
slug: "code",
306+
name: "Code Mode",
307+
roleDefinition: "Test role definition",
308+
groups: ["command", "read", "edit"],
309+
}
310+
})
311+
312+
await newTaskTool(
313+
mockCline as any,
314+
block,
315+
mockAskApproval,
316+
mockHandleError,
317+
mockPushToolResult,
318+
mockRemoveClosingTag,
319+
)
320+
321+
// Verify error was pushed
322+
expect(mockPushToolResult).toHaveBeenCalledWith("Tool Error: Invalid mode: invalid-mode")
323+
324+
// Verify no task was created
325+
expect(mockInitClineWithTask).not.toHaveBeenCalled()
326+
})
187327
})

‎src/core/tools/newTaskTool.ts‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,17 @@ export async function newTaskTool(
7575
return
7676
}
7777

78+
// Check if user selected a different mode during approval
79+
const selectedMode = cline.getAskResponseValues?.selectedMode as string | undefined
80+
const finalMode = selectedMode || mode
81+
82+
// Verify the final mode exists
83+
const finalTargetMode = getModeBySlug(finalMode, (await provider.getState())?.customModes)
84+
if (!finalTargetMode) {
85+
pushToolResult(formatResponse.toolError(`Invalid mode: ${finalMode}`))
86+
return
87+
}
88+
7889
if (cline.enableCheckpoints) {
7990
cline.checkpointSave(true)
8091
}
@@ -89,15 +100,17 @@ export async function newTaskTool(
89100
return
90101
}
91102

92-
// Now switch the newly created task to the desired mode
93-
await provider.handleModeSwitch(mode)
103+
// Now switch the newly created task to the desired mode (using the final mode)
104+
await provider.handleModeSwitch(finalMode)
94105

95106
// Delay to allow mode change to take effect
96107
await delay(500)
97108

98109
cline.emit(RooCodeEventName.TaskSpawned, newCline.taskId)
99110

100-
pushToolResult(`Successfully created new task in ${targetMode.name} mode with message: ${unescapedMessage}`)
111+
pushToolResult(
112+
`Successfully created new task in ${finalTargetMode.name} mode with message: ${unescapedMessage}`,
113+
)
101114

102115
// Set the isPaused flag to true so the parent
103116
// task can wait for the sub-task to finish.

‎webview-ui/src/components/chat/ChatRow.tsx‎

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import MarkdownBlock from "../common/MarkdownBlock"
3333
import { ReasoningBlock } from "./ReasoningBlock"
3434
import Thumbnails from "../common/Thumbnails"
3535
import McpResourceRow from "../mcp/McpResourceRow"
36+
import ModeSelector from "./ModeSelector"
3637

3738
import { Mention } from "./Mention"
3839
import { CheckpointSaved } from "./checkpoints/CheckpointSaved"
@@ -60,6 +61,7 @@ interface ChatRowProps {
6061
onFollowUpUnmount?: () => void
6162
isFollowUpAnswered?: boolean
6263
editable?: boolean
64+
onNewTaskModeChange?: (mode: string) => void
6365
}
6466

6567
// eslint-disable-next-line @typescript-eslint/no-empty-object-type
@@ -112,16 +114,18 @@ export const ChatRowContent = ({
112114
onBatchFileResponse,
113115
isFollowUpAnswered,
114116
editable,
117+
onNewTaskModeChange,
115118
}: ChatRowContentProps) => {
116119
const { t } = useTranslation()
117-
const { mcpServers, alwaysAllowMcp, currentCheckpoint, mode } = useExtensionState()
120+
const { mcpServers, alwaysAllowMcp, currentCheckpoint, mode, customModes, customModePrompts } = useExtensionState()
118121
const [reasoningCollapsed, setReasoningCollapsed] = useState(true)
119122
const [isDiffErrorExpanded, setIsDiffErrorExpanded] = useState(false)
120123
const [showCopySuccess, setShowCopySuccess] = useState(false)
121124
const [isEditing, setIsEditing] = useState(false)
122125
const [editedContent, setEditedContent] = useState("")
123126
const [editMode, setEditMode] = useState<Mode>(mode || "code")
124127
const [editImages, setEditImages] = useState<string[]>([])
128+
const [selectedNewTaskMode, setSelectedNewTaskMode] = useState<Mode | null>(null)
125129
const { copyWithFeedback } = useCopyToClipboard()
126130

127131
// Handle message events for image selection during edit mode
@@ -765,16 +769,39 @@ export const ChatRowContent = ({
765769
</>
766770
)
767771
case "newTask":
772+
// Use the selected mode if available, otherwise use the tool's mode
773+
const effectiveMode = selectedNewTaskMode || tool.mode
774+
768775
return (
769776
<>
770777
<div style={headerStyle}>
771778
{toolIcon("tasklist")}
772-
<span style={{ fontWeight: "bold" }}>
773-
<Trans
774-
i18nKey="chat:subtasks.wantsToCreate"
775-
components={{ code: <code>{tool.mode}</code> }}
776-
values={{ mode: tool.mode }}
777-
/>
779+
<span style={{ fontWeight: "bold", display: "flex", alignItems: "center", gap: "4px" }}>
780+
{t("chat:subtasks.wantsToCreate").split("{mode}")[0]}
781+
{message.type === "ask" ? (
782+
<ModeSelector
783+
value={effectiveMode as Mode}
784+
onChange={(newMode) => {
785+
setSelectedNewTaskMode(newMode)
786+
// Update the tool data with the new mode
787+
if (tool) {
788+
tool.mode = newMode
789+
}
790+
// Notify parent component
791+
onNewTaskModeChange?.(newMode)
792+
}}
793+
disabled={false}
794+
title={t("chat:subtasks.selectMode")}
795+
triggerClassName="inline-flex"
796+
modeShortcutText=""
797+
customModes={customModes}
798+
customModePrompts={customModePrompts}
799+
disableSearch={true}
800+
/>
801+
) : (
802+
<code>{effectiveMode}</code>
803+
)}
804+
{t("chat:subtasks.wantsToCreate").split("{mode}")[1]}
778805
</span>
779806
</div>
780807
<div

‎webview-ui/src/components/chat/ChatView.tsx‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
188188
const autoApproveTimeoutRef = useRef<NodeJS.Timeout | null>(null)
189189
const userRespondedRef = useRef<boolean>(false)
190190
const [currentFollowUpTs, setCurrentFollowUpTs] = useState<number | null>(null)
191+
const [selectedNewTaskModes, setSelectedNewTaskModes] = useState<Record<number, string>>({})
191192

192193
const clineAskRef = useRef(clineAsk)
193194
useEffect(() => {
@@ -709,7 +710,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
709710
// after which buttons are shown and we then send an askResponse to the
710711
// extension.
711712
const handlePrimaryButtonClick = useCallback(
712-
(text?: string, images?: string[]) => {
713+
(text?: string, images?: string[], additionalData?: any) => {
713714
// Mark that user has responded
714715
userRespondedRef.current = true
715716

@@ -730,12 +731,17 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
730731
askResponse: "yesButtonClicked",
731732
text: trimmedInput,
732733
images: images,
734+
values: additionalData,
733735
})
734736
// Clear input state after sending
735737
setInputValue("")
736738
setSelectedImages([])
737739
} else {
738-
vscode.postMessage({ type: "askResponse", askResponse: "yesButtonClicked" })
740+
vscode.postMessage({
741+
type: "askResponse",
742+
askResponse: "yesButtonClicked",
743+
values: additionalData,
744+
})
739745
}
740746
break
741747
case "completion_result":
@@ -1506,6 +1512,9 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
15061512
onBatchFileResponse={handleBatchFileResponse}
15071513
onFollowUpUnmount={handleFollowUpUnmount}
15081514
isFollowUpAnswered={messageOrGroup.ts === currentFollowUpTs}
1515+
onNewTaskModeChange={(mode: string) => {
1516+
setSelectedNewTaskModes((prev) => ({ ...prev, [messageOrGroup.ts]: mode }))
1517+
}}
15091518
editable={
15101519
messageOrGroup.type === "ask" &&
15111520
messageOrGroup.ask === "tool" &&
@@ -1912,7 +1921,26 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
19121921
appearance="primary"
19131922
disabled={!enableButtons}
19141923
className={secondaryButtonText ? "flex-1 mr-[6px]" : "flex-[2] mr-0"}
1915-
onClick={() => handlePrimaryButtonClick(inputValue, selectedImages)}>
1924+
onClick={() => {
1925+
// Check if this is a newTask tool and we have a selected mode
1926+
let additionalData = undefined
1927+
if (lastMessage?.ask === "tool") {
1928+
try {
1929+
const tool = JSON.parse(lastMessage.text || "{}")
1930+
if (
1931+
tool.tool === "newTask" &&
1932+
selectedNewTaskModes[lastMessage.ts]
1933+
) {
1934+
additionalData = {
1935+
selectedMode: selectedNewTaskModes[lastMessage.ts],
1936+
}
1937+
}
1938+
} catch (_e) {
1939+
// Ignore parse errors
1940+
}
1941+
}
1942+
handlePrimaryButtonClick(inputValue, selectedImages, additionalData)
1943+
}}>
19161944
{primaryButtonText}
19171945
</VSCodeButton>
19181946
</StandardTooltip>

‎webview-ui/src/i18n/locales/en/chat.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,8 @@
247247
"completionContent": "Subtask Completed",
248248
"resultContent": "Subtask Results",
249249
"defaultResult": "Please continue to the next task.",
250-
"completionInstructions": "Subtask completed! You can review the results and suggest any corrections or next steps. If everything looks good, confirm to return the result to the parent task."
250+
"completionInstructions": "Subtask completed! You can review the results and suggest any corrections or next steps. If everything looks good, confirm to return the result to the parent task.",
251+
"selectMode": "Click to select a different mode for this subtask"
251252
},
252253
"questions": {
253254
"hasQuestion": "Roo has a question:"

0 commit comments

Comments
 (0)