ui: add upgrade handler for task history migration

Implement a structured upgrade system that manages the task history migration process:

- Create a dedicated upgrade UI that blocks normal app usage until migration completes
- Separate migration check from migration execution for better control flow
- Add progress logging during migration to provide user feedback
- Remove automatic migration during extension activation
- Add new message types for upgrade status and completion

This change improves the user experience during task history migration by providing
visual feedback and ensuring the app is in a consistent state before allowing normal usage.
The upgrade system is designed to be extensible for future structural upgrades beyond
task history migration.

Signed-off-by: Eric Wheeler <[email protected]>

Author: Eric Wheeler

src/core/upgrade/upgrade.ts

@@ -0,0 +1,36 @@
+import { isTaskHistoryMigrationNeeded, migrateTaskHistoryStorage } from "../task-persistence/taskHistory"
+
+/**
+ * Checks if any upgrades are needed in the system.
+ * Currently checks for task history migration needs.
+ *
+ * @returns A promise that resolves to true if any upgrades are needed, false otherwise.
+ */
+export async function isUpgradeNeeded(): Promise<boolean> {
+	// Check if task history migration is needed
+	const taskHistoryMigrationNeeded = await isTaskHistoryMigrationNeeded()
+
+	// Return true if any upgrade is needed
+	return taskHistoryMigrationNeeded
+}
+
+/**
+ * Performs all necessary upgrades in the system.
+ * Currently handles task history migration.
+ *
+ * @param logs Optional array to capture log messages
+ * @returns A promise that resolves to true if upgrades were performed, false if no upgrades were needed.
+ */
+export async function performUpgrade(logs: string[] = []): Promise<boolean> {
+	// Check if task history migration is needed
+	const taskHistoryMigrationNeeded = await isTaskHistoryMigrationNeeded()
+
+	// Perform task history migration if needed
+	if (taskHistoryMigrationNeeded) {
+		await migrateTaskHistoryStorage(logs)
+		return true
+	}
+
+	// No upgrades were needed
+	return false
+}

src/core/webview/webviewMessageHandler.ts

@@ -16,6 +16,7 @@ import {
 	HistoryItem,
 } from "@roo-code/types"
 import { getHistoryItemsForSearch } from "../task-persistence/taskHistory"
+import { isUpgradeNeeded, performUpgrade } from "../upgrade/upgrade"
 import { CloudService } from "@roo-code/cloud"
 import { TelemetryService } from "@roo-code/telemetry"
 import { type ApiMessage } from "../task-persistence/apiMessages"
@@ -67,6 +68,66 @@ export const webviewMessageHandler = async (
 	const getGlobalState = <K extends keyof GlobalState>(key: K) => provider.contextProxy.getValue(key)
 	const updateGlobalState = async <K extends keyof GlobalState>(key: K, value: GlobalState[K]) =>
 		await provider.contextProxy.setValue(key, value)
+	
+	/**
+	 * Helper function to handle common functionality for task history operations
+	 * @param operationName Name of the operation for logging
+	 * @param options Options for the operation
+	 * @param operation The async function to perform
+	 * @param onSuccess Callback for successful operation
+	 * @param onError Callback for operation error
+	 * @param logMessageType Type of message to use when sending logs to UI
+	 */
+	async function handleLoggingOperation<T>(
+		operationName: string,
+		options: any,
+		operation: (options: any, logs: string[]) => Promise<T>,
+		onSuccess: (result: T) => Promise<void>,
+		onError: (error: any) => Promise<void>,
+		logMessageType: "loggingOperation",
+	): Promise<void> {
+		try {
+			// Create a logs array to capture messages
+			const logs: string[] = []
+
+			// Log the options for debugging
+			console.log(`[webviewMessageHandler] ${operationName} options:`, JSON.stringify(options, null, 2))
+
+			// Create a monitoring function to send logs to UI
+			const sendLogsToUI = () => {
+				if (logs.length > 0) {
+					const logsCopy = [...logs]
+					logs.length = 0 // Clear the array
+
+					// Send each log message to the webview
+					for (const log of logsCopy) {
+						provider.postMessageToWebview({
+							type: logMessageType,
+							log,
+						})
+					}
+				}
+			}
+
+			// Set up interval to forward logs during operation
+			const logInterval = setInterval(sendLogsToUI, 100)
+
+			// Perform the operation
+			const result = await operation(options, logs)
+
+			// Clear the interval
+			clearInterval(logInterval)
+
+			// Send any remaining logs
+			sendLogsToUI()
+
+			// Handle success
+			await onSuccess(result)
+		} catch (error) {
+			// Handle error
+			await onError(error)
+		}
+	}
 
 	/**
 	 * Shared utility to find message indices based on timestamp
@@ -2225,6 +2286,65 @@ export const webviewMessageHandler = async (
 			break
 		}
 
+		case "isUpgradeNeeded": {
+			try {
+				const needed = await isUpgradeNeeded()
+				provider.postMessageToWebview({
+					type: "upgradeStatus" as any,
+					values: {
+						needed,
+					},
+				})
+			} catch (error) {
+				console.error(`[Upgrade] webviewMessageHandler: Error in isUpgradeNeeded:`, error)
+				provider.postMessageToWebview({
+					type: "upgradeStatus" as any,
+					values: {
+						needed: false,
+					},
+				})
+			}
+			break
+		}
+
+		case "performUpgrade": {
+			await handleLoggingOperation<{ success: boolean }>(
+				"performUpgrade",
+				{},
+				async (_, logs) => {
+					return { success: await performUpgrade(logs) }
+				},
+				async (result) => {
+					// Then send upgradeComplete message
+					provider.postMessageToWebview({
+						type: "upgradeComplete" as any,
+						values: {
+							success: result.success,
+						},
+					})
+
+					// Finally, send upgradeStatus with needed=false to indicate upgrade is no longer needed
+					provider.postMessageToWebview({
+						type: "upgradeStatus" as any,
+						values: {
+							needed: false,
+						},
+					})
+				},
+				async (error) => {
+					provider.postMessageToWebview({
+						type: "upgradeComplete" as any,
+						values: {
+							success: false,
+							error: String(error),
+						},
+					})
+				},
+				"loggingOperation",
+			)
+			break
+		}
+
 		case "switchTab": {
 			if (message.tab) {
 				// Capture tab shown event for all switchTab messages (which are user-initiated)

src/extension.ts

@@ -71,10 +71,6 @@ export async function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(outputChannel)
 	outputChannel.appendLine(`${Package.name} extension activated - ${JSON.stringify(Package)}`)
 
-	// Initialize and migrate task history storage
-	// (migrateTaskHistoryStorage also calls initializeTaskHistory internally)
-	await migrateTaskHistoryStorage()
-
 	// Migrate old settings to new
 	await migrateSettings(context, outputChannel)
 

src/shared/ExtensionMessage.ts

@@ -107,6 +107,9 @@ export interface ExtensionMessage {
 		| "codeIndexSettingsSaved"
 		| "codeIndexSecretStatus"
 		| "taskDeletedConfirmation"
+		| "loggingOperation"
+		| "upgradeStatus"
+		| "upgradeComplete"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	action?:
@@ -151,6 +154,7 @@ export interface ExtensionMessage {
 	value?: any
 	hasContent?: boolean // For checkRulesDirectoryResult
 	items?: MarketplaceItem[] | HistoryItem[]
+	log?: string
 	userInfo?: CloudUserInfo
 	organizationAllowList?: OrganizationAllowList
 	tab?: string

src/shared/WebviewMessage.ts

@@ -30,6 +30,8 @@ export interface WebviewMessage {
 		| "currentApiConfigName"
 		| "saveApiConfiguration"
 		| "getHistoryItems"
+		| "isUpgradeNeeded"
+		| "performUpgrade"
 		| "upsertApiConfiguration"
 		| "deleteApiConfiguration"
 		| "loadApiConfiguration"

webview-ui/src/App.tsx

@@ -3,6 +3,7 @@ import { useEvent } from "react-use"
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
 
 import { ExtensionMessage } from "@roo/ExtensionMessage"
+import UpgradeHandler, { useUpgradeCheck } from "./components/upgrade/UpgradeHandler"
 import TranslationProvider from "./i18n/TranslationContext"
 import { MarketplaceViewStateManager } from "./components/marketplace/MarketplaceViewStateManager"
 
@@ -55,6 +56,7 @@ const App = () => {
 
 	const [showAnnouncement, setShowAnnouncement] = useState(false)
 	const [tab, setTab] = useState<Tab>("chat")
+	const { upgradeNeeded, clearUpgradeNeeded } = useUpgradeCheck()
 
 	const [humanRelayDialogState, setHumanRelayDialogState] = useState<{
 		isOpen: boolean
@@ -168,6 +170,20 @@ const App = () => {
 
 	// Do not conditionally load ChatView, it's expensive and there's state we
 	// don't want to lose (user input, disableInput, askResponse promise, etc.)
+
+	// Return early while checking for an upgrade because
+	// there may be structures that should not be accessed
+	// until the upgrade completes.
+	if (upgradeNeeded === null) {
+		return null
+	}
+
+	// If an upgrade is needed, show the upgrade UI
+	if (upgradeNeeded) {
+		return <UpgradeHandler onComplete={clearUpgradeNeeded} />
+	}
+
+	// Normal rendering when no upgrade is needed
 	return showWelcome ? (
 		<WelcomeView />
 	) : (

webview-ui/src/__tests__/App.spec.tsx

@@ -86,6 +86,15 @@ vi.mock("@src/components/account/AccountView", () => ({
 	},
 }))
 
+vi.mock("@src/components/upgrade/UpgradeHandler", () => ({
+	__esModule: true,
+	default: () => <div data-testid="upgrade-handler" />,
+	useUpgradeCheck: () => ({
+		upgradeNeeded: false,
+		clearUpgradeNeeded: vi.fn(),
+	}),
+}))
+
 const mockUseExtensionState = vi.fn()
 
 vi.mock("@src/context/ExtensionStateContext", () => ({