feat: make task mode sticky to task (#6177)

Co-authored-by: Roo Code <roomote@roocode.com>
Co-authored-by: hannesrudolph <hrudolph@gmail.com>
Co-authored-by: Daniel Riccio <ricciodaniel98@gmail.com>
Co-authored-by: Daniel <57051444+daniel-lxs@users.noreply.github.com>

packages/types/src/history.ts

@@ -16,6 +16,7 @@ export const historyItemSchema = z.object({
 	totalCost: z.number(),
 	size: z.number().optional(),
 	workspace: z.string().optional(),
+	mode: z.string().optional(),
 })
 
 export type HistoryItem = z.infer<typeof historyItemSchema>

src/core/task-persistence/taskMetadata.ts

@@ -18,6 +18,7 @@ export type TaskMetadataOptions = {
 	taskNumber: number
 	globalStoragePath: string
 	workspace: string
+	mode?: string
 }
 
 export async function taskMetadata({
@@ -26,6 +27,7 @@ export async function taskMetadata({
 	taskNumber,
 	globalStoragePath,
 	workspace,
+	mode,
 }: TaskMetadataOptions) {
 	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
 
@@ -92,6 +94,7 @@ export async function taskMetadata({
 		totalCost: tokenUsage.totalCost,
 		size: taskDirSize,
 		workspace,
+		mode,
 	}
 
 	return { historyItem, tokenUsage }

src/core/task/Task.ts

@@ -137,6 +137,49 @@ export class Task extends EventEmitter<ClineEvents> {
 	readonly parentTask: Task | undefined = undefined
 	readonly taskNumber: number
 	readonly workspacePath: string
+	/**
+	 * The mode associated with this task. Persisted across sessions
+	 * to maintain user context when reopening tasks from history.
+	 *
+	 * ## Lifecycle
+	 *
+	 * ### For new tasks:
+	 * 1. Initially `undefined` during construction
+	 * 2. Asynchronously initialized from provider state via `initializeTaskMode()`
+	 * 3. Falls back to `defaultModeSlug` if provider state is unavailable
+	 *
+	 * ### For history items:
+	 * 1. Immediately set from `historyItem.mode` during construction
+	 * 2. Falls back to `defaultModeSlug` if mode is not stored in history
+	 *
+	 * ## Important
+	 * This property should NOT be accessed directly until `taskModeReady` promise resolves.
+	 * Use `getTaskMode()` for async access or `taskMode` getter for sync access after initialization.
+	 *
+	 * @private
+	 * @see {@link getTaskMode} - For safe async access
+	 * @see {@link taskMode} - For sync access after initialization
+	 * @see {@link waitForModeInitialization} - To ensure initialization is complete
+	 */
+	private _taskMode: string | undefined
+
+	/**
+	 * Promise that resolves when the task mode has been initialized.
+	 * This ensures async mode initialization completes before the task is used.
+	 *
+	 * ## Purpose
+	 * - Prevents race conditions when accessing task mode
+	 * - Ensures provider state is properly loaded before mode-dependent operations
+	 * - Provides a synchronization point for async initialization
+	 *
+	 * ## Resolution timing
+	 * - For history items: Resolves immediately (sync initialization)
+	 * - For new tasks: Resolves after provider state is fetched (async initialization)
+	 *
+	 * @private
+	 * @see {@link waitForModeInitialization} - Public method to await this promise
+	 */
+	private taskModeReady: Promise<void>
 
 	providerRef: WeakRef<ClineProvider>
 	private readonly globalStoragePath: string
@@ -268,9 +311,16 @@ export class Task extends EventEmitter<ClineEvents> {
 		this.parentTask = parentTask
 		this.taskNumber = taskNumber
 
+		// Store the task's mode when it's created
+		// For history items, use the stored mode; for new tasks, we'll set it after getting state
 		if (historyItem) {
+			this._taskMode = historyItem.mode || defaultModeSlug
+			this.taskModeReady = Promise.resolve()
 			TelemetryService.instance.captureTaskRestarted(this.taskId)
 		} else {
+			// For new tasks, don't set the mode yet - wait for async initialization
+			this._taskMode = undefined
+			this.taskModeReady = this.initializeTaskMode(provider)
 			TelemetryService.instance.captureTaskCreated(this.taskId)
 		}
 
@@ -307,6 +357,129 @@ export class Task extends EventEmitter<ClineEvents> {
 		}
 	}
 
+	/**
+	 * Initialize the task mode from the provider state.
+	 * This method handles async initialization with proper error handling.
+	 *
+	 * ## Flow
+	 * 1. Attempts to fetch the current mode from provider state
+	 * 2. Sets `_taskMode` to the fetched mode or `defaultModeSlug` if unavailable
+	 * 3. Handles errors gracefully by falling back to default mode
+	 * 4. Logs any initialization errors for debugging
+	 *
+	 * ## Error handling
+	 * - Network failures when fetching provider state
+	 * - Provider not yet initialized
+	 * - Invalid state structure
+	 *
+	 * All errors result in fallback to `defaultModeSlug` to ensure task can proceed.
+	 *
+	 * @private
+	 * @param provider - The ClineProvider instance to fetch state from
+	 * @returns Promise that resolves when initialization is complete
+	 */
+	private async initializeTaskMode(provider: ClineProvider): Promise<void> {
+		try {
+			const state = await provider.getState()
+			this._taskMode = state?.mode || defaultModeSlug
+		} catch (error) {
+			// If there's an error getting state, use the default mode
+			this._taskMode = defaultModeSlug
+			// Use the provider's log method for better error visibility
+			const errorMessage = `Failed to initialize task mode: ${error instanceof Error ? error.message : String(error)}`
+			provider.log(errorMessage)
+		}
+	}
+
+	/**
+	 * Wait for the task mode to be initialized before proceeding.
+	 * This method ensures that any operations depending on the task mode
+	 * will have access to the correct mode value.
+	 *
+	 * ## When to use
+	 * - Before accessing mode-specific configurations
+	 * - When switching between tasks with different modes
+	 * - Before operations that depend on mode-based permissions
+	 *
+	 * ## Example usage
+	 * ```typescript
+	 * // Wait for mode initialization before mode-dependent operations
+	 * await task.waitForModeInitialization();
+	 * const mode = task.taskMode; // Now safe to access synchronously
+	 *
+	 * // Or use with getTaskMode() for a one-liner
+	 * const mode = await task.getTaskMode(); // Internally waits for initialization
+	 * ```
+	 *
+	 * @returns Promise that resolves when the task mode is initialized
+	 * @public
+	 */
+	public async waitForModeInitialization(): Promise<void> {
+		return this.taskModeReady
+	}
+
+	/**
+	 * Get the task mode asynchronously, ensuring it's properly initialized.
+	 * This is the recommended way to access the task mode as it guarantees
+	 * the mode is available before returning.
+	 *
+	 * ## Async behavior
+	 * - Internally waits for `taskModeReady` promise to resolve
+	 * - Returns the initialized mode or `defaultModeSlug` as fallback
+	 * - Safe to call multiple times - subsequent calls return immediately if already initialized
+	 *
+	 * ## Example usage
+	 * ```typescript
+	 * // Safe async access
+	 * const mode = await task.getTaskMode();
+	 * console.log(`Task is running in ${mode} mode`);
+	 *
+	 * // Use in conditional logic
+	 * if (await task.getTaskMode() === 'architect') {
+	 *   // Perform architect-specific operations
+	 * }
+	 * ```
+	 *
+	 * @returns Promise resolving to the task mode string
+	 * @public
+	 */
+	public async getTaskMode(): Promise<string> {
+		await this.taskModeReady
+		return this._taskMode || defaultModeSlug
+	}
+
+	/**
+	 * Get the task mode synchronously. This should only be used when you're certain
+	 * that the mode has already been initialized (e.g., after waitForModeInitialization).
+	 *
+	 * ## When to use
+	 * - In synchronous contexts where async/await is not available
+	 * - After explicitly waiting for initialization via `waitForModeInitialization()`
+	 * - In event handlers or callbacks where mode is guaranteed to be initialized
+	 *
+	 * ## Example usage
+	 * ```typescript
+	 * // After ensuring initialization
+	 * await task.waitForModeInitialization();
+	 * const mode = task.taskMode; // Safe synchronous access
+	 *
+	 * // In an event handler after task is started
+	 * task.on('taskStarted', () => {
+	 *   console.log(`Task started in ${task.taskMode} mode`); // Safe here
+	 * });
+	 * ```
+	 *
+	 * @throws {Error} If the mode hasn't been initialized yet
+	 * @returns The task mode string
+	 * @public
+	 */
+	public get taskMode(): string {
+		if (this._taskMode === undefined) {
+			throw new Error("Task mode accessed before initialization. Use getTaskMode() or wait for taskModeReady.")
+		}
+		return this._taskMode
+	}
+
 	static create(options: TaskOptions): [Task, Promise<void>] {
 		const instance = new Task({ ...options, startTask: false })
 		const { images, task, historyItem } = options
@@ -411,6 +584,7 @@ export class Task extends EventEmitter<ClineEvents> {
 				taskNumber: this.taskNumber,
 				globalStoragePath: this.globalStoragePath,
 				workspace: this.cwd,
+				mode: this._taskMode || defaultModeSlug, // Use the task's own mode, not the current provider mode
 			})
 
 			this.emit("taskTokenUsageUpdated", this.taskId, tokenUsage)

src/core/tools/newTaskTool.ts

@@ -80,17 +80,19 @@ export async function newTaskTool(
 			// Preserve the current mode so we can resume with it later.
 			cline.pausedModeSlug = (await provider.getState()).mode ?? defaultModeSlug
 
-			// Switch mode first, then create new task instance.
-			await provider.handleModeSwitch(mode)
-
-			// Delay to allow mode change to take effect before next tool is executed.
-			await delay(500)
-
+			// Create new task instance first (this preserves parent's current mode in its history)
 			const newCline = await provider.initClineWithTask(unescapedMessage, undefined, cline)
 			if (!newCline) {
 				pushToolResult(t("tools:newTask.errors.policy_restriction"))
 				return
 			}
+
+			// Now switch the newly created task to the desired mode
+			await provider.handleModeSwitch(mode)
+
+			// Delay to allow mode change to take effect
+			await delay(500)
+
 			cline.emit("taskSpawned", newCline.taskId)
 
 			pushToolResult(`Successfully created new task in ${targetMode.name} mode with message: ${unescapedMessage}`)

src/core/webview/ClineProvider.ts

@@ -40,7 +40,7 @@ import { findLast } from "../../shared/array"
 import { supportPrompt } from "../../shared/support-prompt"
 import { GlobalFileNames } from "../../shared/globalFileNames"
 import { ExtensionMessage, MarketplaceInstalledMetadata } from "../../shared/ExtensionMessage"
-import { Mode, defaultModeSlug } from "../../shared/modes"
+import { Mode, defaultModeSlug, getModeBySlug } from "../../shared/modes"
 import { experimentDefault, experiments, EXPERIMENT_IDS } from "../../shared/experiments"
 import { formatLanguage } from "../../shared/language"
 import { DEFAULT_WRITE_DELAY_MS } from "@roo-code/types"
@@ -578,6 +578,49 @@ export class ClineProvider
 	public async initClineWithHistoryItem(historyItem: HistoryItem & { rootTask?: Task; parentTask?: Task }) {
 		await this.removeClineFromStack()
 
+		// If the history item has a saved mode, restore it and its associated API configuration
+		if (historyItem.mode) {
+			// Validate that the mode still exists
+			const customModes = await this.customModesManager.getCustomModes()
+			const modeExists = getModeBySlug(historyItem.mode, customModes) !== undefined
+
+			if (!modeExists) {
+				// Mode no longer exists, fall back to default mode
+				this.log(
+					`Mode '${historyItem.mode}' from history no longer exists. Falling back to default mode '${defaultModeSlug}'.`,
+				)
+				historyItem.mode = defaultModeSlug
+			}
+
+			await this.updateGlobalState("mode", historyItem.mode)
+
+			// Load the saved API config for the restored mode if it exists
+			const savedConfigId = await this.providerSettingsManager.getModeConfigId(historyItem.mode)
+			const listApiConfig = await this.providerSettingsManager.listConfig()
+
+			// Update listApiConfigMeta first to ensure UI has latest data
+			await this.updateGlobalState("listApiConfigMeta", listApiConfig)
+
+			// If this mode has a saved config, use it
+			if (savedConfigId) {
+				const profile = listApiConfig.find(({ id }) => id === savedConfigId)
+
+				if (profile?.name) {
+					try {
+						await this.activateProviderProfile({ name: profile.name })
+					} catch (error) {
+						// Log the error but continue with task restoration
+						this.log(
+							`Failed to restore API configuration for mode '${historyItem.mode}': ${
+								error instanceof Error ? error.message : String(error)
+							}. Continuing with default configuration.`,
+						)
+						// The task will continue with the current/default configuration
+					}
+				}
+			}
+		}
+
 		const {
 			apiConfiguration,
 			diffEnabled: enableDiff,
@@ -807,6 +850,31 @@ export class ClineProvider
 		if (cline) {
 			TelemetryService.instance.captureModeSwitch(cline.taskId, newMode)
 			cline.emit("taskModeSwitched", cline.taskId, newMode)
+
+			// Store the current mode in case we need to rollback
+			const previousMode = (cline as any)._taskMode
+
+			try {
+				// Update the task history with the new mode first
+				const history = this.getGlobalState("taskHistory") ?? []
+				const taskHistoryItem = history.find((item) => item.id === cline.taskId)
+				if (taskHistoryItem) {
+					taskHistoryItem.mode = newMode
+					await this.updateTaskHistory(taskHistoryItem)
+				}
+
+				// Only update the task's mode after successful persistence
+				;(cline as any)._taskMode = newMode
+			} catch (error) {
+				// If persistence fails, log the error but don't update the in-memory state
+				this.log(
+					`Failed to persist mode switch for task ${cline.taskId}: ${error instanceof Error ? error.message : String(error)}`,
+				)
+
+				// Optionally, we could emit an event to notify about the failure
+				// This ensures the in-memory state remains consistent with persisted state
+				throw error
+			}
 		}
 
 		await this.updateGlobalState("mode", newMode)

src/core/webview/__tests__/ClineProvider.spec.ts

@@ -1654,6 +1654,268 @@ describe("ClineProvider", () => {
 		})
 	})
 
+	describe("initClineWithHistoryItem mode validation", () => {
+		test("validates and falls back to default mode when restored mode no longer exists", async () => {
+			await provider.resolveWebviewView(mockWebviewView)
+
+			// Mock custom modes that don't include the saved mode
+			const mockCustomModesManager = {
+				getCustomModes: vi.fn().mockResolvedValue([
+					{
+						slug: "existing-mode",
+						name: "Existing Mode",
+						roleDefinition: "Test role",
+						groups: ["read"] as const,
+					},
+				]),
+				dispose: vi.fn(),
+			}
+			;(provider as any).customModesManager = mockCustomModesManager
+
+			// Mock getModeBySlug to return undefined for non-existent mode
+			const { getModeBySlug } = await import("../../../shared/modes")
+			vi.mocked(getModeBySlug)
+				.mockReturnValueOnce(undefined) // First call returns undefined (mode doesn't exist)
+				.mockReturnValue({
+					slug: "code",
+					name: "Code Mode",
+					roleDefinition: "You are a code assistant",
+					groups: ["read", "edit", "browser"],
+				}) // Subsequent calls return default mode
+
+			// Mock provider settings manager
+			;(provider as any).providerSettingsManager = {
+				getModeConfigId: vi.fn().mockResolvedValue(undefined),
+				listConfig: vi.fn().mockResolvedValue([]),
+			}
+
+			// Spy on log method to verify warning was logged
+			const logSpy = vi.spyOn(provider, "log")
+
+			// Create history item with non-existent mode
+			const historyItem = {
+				id: "test-id",
+				ts: Date.now(),
+				task: "Test task",
+				mode: "non-existent-mode", // This mode doesn't exist
+				number: 1,
+				tokensIn: 0,
+				tokensOut: 0,
+				totalCost: 0,
+			}
+
+			// Initialize with history item
+			await provider.initClineWithHistoryItem(historyItem)
+
+			// Verify mode validation occurred
+			expect(mockCustomModesManager.getCustomModes).toHaveBeenCalled()
+			expect(getModeBySlug).toHaveBeenCalledWith("non-existent-mode", expect.any(Array))
+
+			// Verify fallback to default mode
+			expect(mockContext.globalState.update).toHaveBeenCalledWith("mode", "code")
+			expect(logSpy).toHaveBeenCalledWith(
+				"Mode 'non-existent-mode' from history no longer exists. Falling back to default mode 'code'.",
+			)
+
+			// Verify history item was updated with default mode
+			expect(historyItem.mode).toBe("code")
+		})
+
+		test("preserves mode when it exists in custom modes", async () => {
+			await provider.resolveWebviewView(mockWebviewView)
+
+			// Mock custom modes that include the saved mode
+			const mockCustomModesManager = {
+				getCustomModes: vi.fn().mockResolvedValue([
+					{
+						slug: "custom-mode",
+						name: "Custom Mode",
+						roleDefinition: "Custom role",
+						groups: ["read", "edit"] as const,
+					},
+				]),
+				dispose: vi.fn(),
+			}
+			;(provider as any).customModesManager = mockCustomModesManager
+
+			// Mock getModeBySlug to return the custom mode
+			const { getModeBySlug } = await import("../../../shared/modes")
+			vi.mocked(getModeBySlug).mockReturnValue({
+				slug: "custom-mode",
+				name: "Custom Mode",
+				roleDefinition: "Custom role",
+				groups: ["read", "edit"],
+			})
+
+			// Mock provider settings manager
+			;(provider as any).providerSettingsManager = {
+				getModeConfigId: vi.fn().mockResolvedValue("config-id"),
+				listConfig: vi
+					.fn()
+					.mockResolvedValue([{ name: "test-config", id: "config-id", apiProvider: "anthropic" }]),
+				activateProfile: vi
+					.fn()
+					.mockResolvedValue({ name: "test-config", id: "config-id", apiProvider: "anthropic" }),
+			}
+
+			// Spy on log method to verify no warning was logged
+			const logSpy = vi.spyOn(provider, "log")
+
+			// Create history item with existing custom mode
+			const historyItem = {
+				id: "test-id",
+				ts: Date.now(),
+				task: "Test task",
+				mode: "custom-mode",
+				number: 1,
+				tokensIn: 0,
+				tokensOut: 0,
+				totalCost: 0,
+			}
+
+			// Initialize with history item
+			await provider.initClineWithHistoryItem(historyItem)
+
+			// Verify mode validation occurred
+			expect(mockCustomModesManager.getCustomModes).toHaveBeenCalled()
+			expect(getModeBySlug).toHaveBeenCalledWith("custom-mode", expect.any(Array))
+
+			// Verify mode was preserved
+			expect(mockContext.globalState.update).toHaveBeenCalledWith("mode", "custom-mode")
+			expect(logSpy).not.toHaveBeenCalledWith(expect.stringContaining("no longer exists"))
+
+			// Verify history item mode was not changed
+			expect(historyItem.mode).toBe("custom-mode")
+		})
+
+		test("preserves mode when it exists in built-in modes", async () => {
+			await provider.resolveWebviewView(mockWebviewView)
+
+			// Mock no custom modes
+			const mockCustomModesManager = {
+				getCustomModes: vi.fn().mockResolvedValue([]),
+				dispose: vi.fn(),
+			}
+			;(provider as any).customModesManager = mockCustomModesManager
+
+			// Mock getModeBySlug to return built-in architect mode
+			const { getModeBySlug } = await import("../../../shared/modes")
+			vi.mocked(getModeBySlug).mockReturnValue({
+				slug: "architect",
+				name: "Architect Mode",
+				roleDefinition: "You are an architect",
+				groups: ["read", "edit"],
+			})
+
+			// Mock provider settings manager
+			;(provider as any).providerSettingsManager = {
+				getModeConfigId: vi.fn().mockResolvedValue(undefined),
+				listConfig: vi.fn().mockResolvedValue([]),
+			}
+
+			// Create history item with built-in mode
+			const historyItem = {
+				id: "test-id",
+				ts: Date.now(),
+				task: "Test task",
+				mode: "architect",
+				number: 1,
+				tokensIn: 0,
+				tokensOut: 0,
+				totalCost: 0,
+			}
+
+			// Initialize with history item
+			await provider.initClineWithHistoryItem(historyItem)
+
+			// Verify mode was preserved
+			expect(mockContext.globalState.update).toHaveBeenCalledWith("mode", "architect")
+
+			// Verify history item mode was not changed
+			expect(historyItem.mode).toBe("architect")
+		})
+
+		test("handles history items without mode property", async () => {
+			await provider.resolveWebviewView(mockWebviewView)
+
+			// Mock provider settings manager
+			;(provider as any).providerSettingsManager = {
+				getModeConfigId: vi.fn().mockResolvedValue(undefined),
+				listConfig: vi.fn().mockResolvedValue([]),
+			}
+
+			// Create history item without mode
+			const historyItem = {
+				id: "test-id",
+				ts: Date.now(),
+				task: "Test task",
+				// No mode property
+				number: 1,
+				tokensIn: 0,
+				tokensOut: 0,
+				totalCost: 0,
+			}
+
+			// Initialize with history item
+			await provider.initClineWithHistoryItem(historyItem)
+
+			// Verify no mode validation occurred (mode update not called)
+			expect(mockContext.globalState.update).not.toHaveBeenCalledWith("mode", expect.any(String))
+		})
+
+		test("continues with task restoration even if mode config loading fails", async () => {
+			await provider.resolveWebviewView(mockWebviewView)
+
+			// Mock custom modes
+			const mockCustomModesManager = {
+				getCustomModes: vi.fn().mockResolvedValue([]),
+				dispose: vi.fn(),
+			}
+			;(provider as any).customModesManager = mockCustomModesManager
+
+			// Mock getModeBySlug to return built-in mode
+			const { getModeBySlug } = await import("../../../shared/modes")
+			vi.mocked(getModeBySlug).mockReturnValue({
+				slug: "code",
+				name: "Code Mode",
+				roleDefinition: "You are a code assistant",
+				groups: ["read", "edit", "browser"],
+			})
+
+			// Mock provider settings manager to throw error
+			;(provider as any).providerSettingsManager = {
+				getModeConfigId: vi.fn().mockResolvedValue("config-id"),
+				listConfig: vi
+					.fn()
+					.mockResolvedValue([{ name: "test-config", id: "config-id", apiProvider: "anthropic" }]),
+				activateProfile: vi.fn().mockRejectedValue(new Error("Failed to load config")),
+			}
+
+			// Spy on log method
+			const logSpy = vi.spyOn(provider, "log")
+
+			// Create history item
+			const historyItem = {
+				id: "test-id",
+				ts: Date.now(),
+				task: "Test task",
+				mode: "code",
+				number: 1,
+				tokensIn: 0,
+				tokensOut: 0,
+				totalCost: 0,
+			}
+
+			// Initialize with history item - should not throw
+			await expect(provider.initClineWithHistoryItem(historyItem)).resolves.not.toThrow()
+
+			// Verify error was logged but task restoration continued
+			expect(logSpy).toHaveBeenCalledWith(
+				expect.stringContaining("Failed to restore API configuration for mode 'code'"),
+			)
+		})
+	})
+
 	describe("updateCustomMode", () => {
 		test("updates both file and state when updating custom mode", async () => {
 			await provider.resolveWebviewView(mockWebviewView)