diff --git a/packages/types/src/mode.ts b/packages/types/src/mode.ts index 88dcbb957471..906248f34642 100644 --- a/packages/types/src/mode.ts +++ b/packages/types/src/mode.ts @@ -70,6 +70,13 @@ export const modeConfigSchema = z.object({ customInstructions: z.string().optional(), groups: groupEntryArraySchema, source: z.enum(["global", "project"]).optional(), + onComplete: z + .object({ + switchToMode: z.string().optional(), + runCommand: z.string().optional(), + includeSummary: z.boolean().optional(), + }) + .optional(), }) export type ModeConfig = z.infer diff --git a/src/core/tools/__tests__/attemptCompletionTool.spec.ts b/src/core/tools/__tests__/attemptCompletionTool.spec.ts index fcad4d5f4925..b0d89b7ddb51 100644 --- a/src/core/tools/__tests__/attemptCompletionTool.spec.ts +++ b/src/core/tools/__tests__/attemptCompletionTool.spec.ts @@ -1,4 +1,4 @@ -import { TodoItem } from "@roo-code/types" +import { TodoItem, ModeConfig } from "@roo-code/types" import { AttemptCompletionToolUse } from "../../../shared/tools" @@ -6,6 +6,21 @@ import { AttemptCompletionToolUse } from "../../../shared/tools" vi.mock("../../prompts/responses", () => ({ formatResponse: { toolError: vi.fn((msg: string) => `Error: ${msg}`), + imageBlocks: vi.fn(() => []), + }, +})) + +// Mock the getModeConfig function +vi.mock("../../../shared/modes", () => ({ + getModeConfig: vi.fn(), +})) + +// Mock TelemetryService +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureTaskCompleted: vi.fn(), + }, }, })) @@ -28,6 +43,7 @@ vi.mock("../../../shared/package", () => ({ import { attemptCompletionTool } from "../attemptCompletionTool" import { Task } from "../../task/Task" import * as vscode from "vscode" +import { getModeConfig } from "../../../shared/modes" describe("attemptCompletionTool", () => { let mockTask: Partial @@ -38,13 +54,21 @@ describe("attemptCompletionTool", () => { let mockToolDescription: ReturnType let mockAskFinishSubTaskApproval: ReturnType let mockGetConfiguration: ReturnType + let mockProvider: any + let mockGetState: ReturnType + let mockHandleModeSwitch: ReturnType + let mockPostMessageToWebview: ReturnType + let mockSay: ReturnType + let mockAsk: ReturnType + let mockEmit: ReturnType + let mockProviderRef: any beforeEach(() => { mockPushToolResult = vi.fn() mockAskApproval = vi.fn() mockHandleError = vi.fn() - mockRemoveClosingTag = vi.fn() - mockToolDescription = vi.fn() + mockRemoveClosingTag = vi.fn((tag, content) => content || "") + mockToolDescription = vi.fn(() => "attempt_completion") mockAskFinishSubTaskApproval = vi.fn() mockGetConfiguration = vi.fn(() => ({ get: vi.fn((key: string, defaultValue: any) => { @@ -58,11 +82,56 @@ describe("attemptCompletionTool", () => { // Setup vscode mock vi.mocked(vscode.workspace.getConfiguration).mockImplementation(mockGetConfiguration) + // Mock provider methods + mockGetState = vi.fn() + mockHandleModeSwitch = vi.fn() + mockPostMessageToWebview = vi.fn() + + mockProvider = { + getState: mockGetState, + handleModeSwitch: mockHandleModeSwitch, + postMessageToWebview: mockPostMessageToWebview, + } + + mockProviderRef = { + deref: vi.fn(() => mockProvider), + } + + mockSay = vi.fn() + mockAsk = vi.fn(() => Promise.resolve({ response: "yesButtonClicked", text: "", images: [] })) + mockEmit = vi.fn() + mockTask = { consecutiveMistakeCount: 0, recordToolError: vi.fn(), todoList: undefined, + providerRef: mockProviderRef, + say: mockSay, + ask: mockAsk, + emit: mockEmit, + clineMessages: [], + userMessageContent: [], + taskId: "test-task-id", + getTokenUsage: vi.fn(() => ({ + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.001, + contextTokens: 80, + totalCacheWrites: 0, + totalCacheReads: 0, + })), + toolUsage: {}, + parentTask: undefined, + sayAndCreateMissingParamError: vi.fn(async (tool, param) => `Missing parameter: ${param} for ${tool}`), } + + // Reset getModeConfig mock + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + } as ModeConfig) }) describe("todo list validation", () => { @@ -409,4 +478,286 @@ describe("attemptCompletionTool", () => { ) }) }) + + describe("onComplete actions", () => { + beforeEach(() => { + // Setup default state for onComplete tests + mockGetState.mockResolvedValue({ + mode: "code", + customModes: [], + }) + }) + + it("should execute mode switch when onComplete.switchToMode is configured", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Configure mode with onComplete.switchToMode + vi.mocked(getModeConfig).mockReturnValue({ + slug: "architect", + name: "Architect", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + switchToMode: "code", + }, + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify mode switch was called + expect(mockHandleModeSwitch).toHaveBeenCalledWith("code") + expect(mockSay).toHaveBeenCalledWith("text", "Automatically switching to code mode as configured...") + }) + + it("should execute command when onComplete.runCommand is configured", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Configure mode with onComplete.runCommand + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + runCommand: "/test-command", + }, + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify command was executed + expect(mockPostMessageToWebview).toHaveBeenCalledWith({ + type: "invoke", + invoke: "sendMessage", + text: "/test-command", + }) + expect(mockSay).toHaveBeenCalledWith("text", "Automatically executing command: /test-command") + }) + + it("should include summary when onComplete.includeSummary is true", async () => { + const taskResult = "Task completed with specific results" + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: taskResult }, + partial: false, + } + + // Configure mode with onComplete.runCommand and includeSummary + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + runCommand: "/review", + includeSummary: true, + }, + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify command was executed with summary + expect(mockPostMessageToWebview).toHaveBeenCalledWith({ + type: "invoke", + invoke: "sendMessage", + text: `/review\n\nContext from previous task:\n${taskResult}`, + }) + }) + + it("should execute both mode switch and command when both are configured", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Configure mode with both onComplete actions + vi.mocked(getModeConfig).mockReturnValue({ + slug: "architect", + name: "Architect", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + switchToMode: "debug", + runCommand: "/analyze", + includeSummary: false, + }, + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify both actions were executed + expect(mockHandleModeSwitch).toHaveBeenCalledWith("debug") + expect(mockPostMessageToWebview).toHaveBeenCalledWith({ + type: "invoke", + invoke: "sendMessage", + text: "/analyze", + }) + }) + + it("should handle errors in onComplete actions gracefully", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Configure mode with onComplete.switchToMode that will fail + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + switchToMode: "invalid-mode", + }, + } as ModeConfig) + + // Make handleModeSwitch throw an error + mockHandleModeSwitch.mockRejectedValue(new Error("Invalid mode")) + + // Add console.error spy + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify error was logged but didn't break completion + expect(consoleErrorSpy).toHaveBeenCalledWith("Failed to execute onComplete actions:", expect.any(Error)) + + // Verify completion still succeeded (ask was called for completion_result) + expect(mockAsk).toHaveBeenCalledWith("completion_result", "", false) + + consoleErrorSpy.mockRestore() + }) + + it("should not execute onComplete actions when mode has no onComplete config", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Mode without onComplete configuration + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + // No onComplete field + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify no onComplete actions were executed + expect(mockHandleModeSwitch).not.toHaveBeenCalled() + expect(mockPostMessageToWebview).not.toHaveBeenCalled() + }) + + it("should not execute onComplete actions when provider is not available", async () => { + const block: AttemptCompletionToolUse = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Task completed successfully" }, + partial: false, + } + + // Make provider unavailable + mockProviderRef.deref.mockReturnValue(null) + + // Configure mode with onComplete actions + vi.mocked(getModeConfig).mockReturnValue({ + slug: "code", + name: "Code", + roleDefinition: "Test role", + groups: ["read", "edit"], + onComplete: { + switchToMode: "architect", + runCommand: "/test", + }, + } as ModeConfig) + + await attemptCompletionTool( + mockTask as Task, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + mockToolDescription, + mockAskFinishSubTaskApproval, + ) + + // Verify no onComplete actions were executed + expect(mockHandleModeSwitch).not.toHaveBeenCalled() + expect(mockPostMessageToWebview).not.toHaveBeenCalled() + }) + }) }) diff --git a/src/core/tools/attemptCompletionTool.ts b/src/core/tools/attemptCompletionTool.ts index 5074d7f4e808..36ff11e3a9f0 100644 --- a/src/core/tools/attemptCompletionTool.ts +++ b/src/core/tools/attemptCompletionTool.ts @@ -17,6 +17,7 @@ import { } from "../../shared/tools" import { formatResponse } from "../prompts/responses" import { Package } from "../../shared/package" +import { getModeConfig } from "../../shared/modes" export async function attemptCompletionTool( cline: Task, @@ -95,6 +96,52 @@ export async function attemptCompletionTool( TelemetryService.instance.captureTaskCompleted(cline.taskId) cline.emit(RooCodeEventName.TaskCompleted, cline.taskId, cline.getTokenUsage(), cline.toolUsage) + // Check for onComplete actions in the current mode configuration + const provider = cline.providerRef.deref() + if (provider) { + const state = await provider.getState() + const currentMode = state?.mode || "code" + const customModes = state?.customModes || [] + + try { + const modeConfig = getModeConfig(currentMode, customModes) + + if (modeConfig.onComplete) { + const { switchToMode, runCommand, includeSummary } = modeConfig.onComplete + + // Prepare context for onComplete actions + const context = includeSummary ? result : undefined + + // Execute mode switch if specified + if (switchToMode) { + await cline.say("text", `Automatically switching to ${switchToMode} mode as configured...`) + await provider.handleModeSwitch(switchToMode) + } + + // Execute command if specified + if (runCommand) { + await cline.say("text", `Automatically executing command: ${runCommand}`) + + // Create a new message with the command, optionally including the summary + let commandMessage = runCommand + if (context) { + commandMessage = `${runCommand}\n\nContext from previous task:\n${context}` + } + + // Send the command as a new user message + await provider.postMessageToWebview({ + type: "invoke", + invoke: "sendMessage", + text: commandMessage, + }) + } + } + } catch (error) { + // Log error but don't fail the completion + console.error("Failed to execute onComplete actions:", error) + } + } + if (cline.parentTask) { const didApprove = await askFinishSubTaskApproval()