[Migration Tool] Implement AI enhancement feature for migration tools

workspaces/ballerina/ballerina-core/src/rpc-types/migrate-integration/index.ts

@@ -17,7 +17,7 @@
  */
 
 import { ImportIntegrationResponse } from "../../interfaces/extended-lang-client";
-import { GetMigrationToolsResponse, ImportIntegrationRPCRequest, MigrateRequest, MigrationToolPullRequest, OpenMigrationReportRequest, OpenSubProjectReportRequest, SaveMigrationReportRequest, StoreSubProjectReportsRequest } from "./interfaces";
+import { ActiveMigrationSession, GetMigrationToolsResponse, ImportIntegrationRPCRequest, MigrateRequest, MigrationToolPullRequest, OpenMigrationReportRequest, OpenSubProjectReportRequest, SaveMigrationReportRequest, StoreSubProjectReportsRequest } from "./interfaces";
 
 export interface MigrateIntegrationAPI {
     getMigrationTools: () => Promise<GetMigrationToolsResponse>;
@@ -28,4 +28,15 @@ export interface MigrateIntegrationAPI {
     storeSubProjectReports: (params: StoreSubProjectReportsRequest) => void;
     saveMigrationReport: (params: SaveMigrationReportRequest) => void;
     migrateProject: (params: MigrateRequest) => void;
+    getActiveMigrationSession: () => Promise<ActiveMigrationSession>;
+    /** Starts (or resumes) the AI enhancement pipeline from AI Chat. */
+    startMigrationEnhancement: () => Promise<void>;
+    /** Seeds any persisted migration conversation history into the chat state. */
+    seedMigrationHistory: () => Promise<boolean>;
+    /** Triggers the wizard-level streaming enhancement pipeline (called from AI Chat after startMigrationEnhancement). */
+    wizardEnhancementReady: () => Promise<void>;
+    /** Aborts the currently running migration AI agent. */
+    abortMigrationAgent: () => Promise<void>;
+    /** Retrieves the persisted migration conversation history messages for display in AI Chat. */
+    getMigrationHistoryMessages: () => Promise<Array<{ role: string; content: string }>>;
 }

workspaces/ballerina/ballerina-core/src/rpc-types/migrate-integration/interfaces.ts

@@ -65,12 +65,29 @@ export interface SaveMigrationReportRequest {
     };
 }
 
+/**
+ * Describes the current state of a migration AI enhancement session.
+ * Returned by `getActiveMigrationSession` so the UI can show a live banner.
+ */
+export interface ActiveMigrationSession {
+    /** `true` while the enhancement pipeline is running. */
+    isActive: boolean;
+    /** `true` when the AI enhancement feature was used (wizard or post-wizard). */
+    aiFeatureUsed: boolean;
+    /** `true` once the AI enhancement pipeline has completed successfully. */
+    fullyEnhanced: boolean;
+}
+
 export interface MigrateRequest {
     project: ProjectRequest;
     textEdits: {
         [key: string]: string;
     };
     projects?: ProjectMigrationResult[];
+    /** `true` when the AI enhancement feature was used (wizard or post-wizard). */
+    aiFeatureUsed?: boolean;
+    /** Absolute path to the original source project (e.g. Mule XML directory). */
+    sourcePath?: string;
 }
 
 export interface OpenSubProjectReportRequest {

workspaces/ballerina/ballerina-core/src/rpc-types/migrate-integration/rpc-type.ts

@@ -18,7 +18,7 @@
  * THIS FILE INCLUDES AUTO GENERATED CODE
  */
 import { ImportIntegrationResponse } from "../../interfaces/extended-lang-client";
-import { GetMigrationToolsResponse, ImportIntegrationRPCRequest, MigrateRequest, MigrationToolPullRequest, OpenMigrationReportRequest, OpenSubProjectReportRequest, SaveMigrationReportRequest, StoreSubProjectReportsRequest } from "./interfaces";
+import { ActiveMigrationSession, GetMigrationToolsResponse, ImportIntegrationRPCRequest, MigrateRequest, MigrationToolPullRequest, OpenMigrationReportRequest, OpenSubProjectReportRequest, SaveMigrationReportRequest, StoreSubProjectReportsRequest } from "./interfaces";
 import { RequestType, NotificationType } from "vscode-messenger-common";
 
 const _preFix = "migrate-integration";
@@ -30,3 +30,5 @@ export const openSubProjectReport: NotificationType<OpenSubProjectReportRequest>
 export const storeSubProjectReports: NotificationType<StoreSubProjectReportsRequest> = { method: `${_preFix}/storeSubProjectReports` };
 export const saveMigrationReport: NotificationType<SaveMigrationReportRequest> = { method: `${_preFix}/saveMigrationReport` };
 export const migrateProject: NotificationType<MigrateRequest> = { method: `${_preFix}/migrateProject` };
+export const getActiveMigrationSession: RequestType<void, ActiveMigrationSession> = { method: `${_preFix}/getActiveMigrationSession` };
+export const seedMigrationHistory: RequestType<void, boolean> = { method: `${_preFix}/seedMigrationHistory` };

workspaces/ballerina/ballerina-extension/src/RPCLayer.ts

@@ -47,6 +47,7 @@ import { activeAgentChanged } from '@wso2/ballerina-core';
 import { ArtifactsUpdated, ArtifactNotificationHandler } from './utils/project-artifacts-handler';
 import { registerMigrateIntegrationRpcHandlers } from './rpc-managers/migrate-integration/rpc-handler';
 import { registerPlatformExtRpcHandlers } from './rpc-managers/platform-ext/rpc-handler';
+import { MigrationPanelWebview } from './views/migration-panel/webview';
 
 export class RPCLayer {
     static _messenger: Messenger = new Messenger();
@@ -64,6 +65,9 @@ export class RPCLayer {
             window.onDidChangeActiveColorTheme((theme) => {
                 RPCLayer._messenger.sendNotification(currentThemeChanged, { type: 'webview', webviewType: VisualizerWebview.viewType }, theme.kind);
             });
+        } else if (isMigrationPanel(webViewPanel)) {
+            // Migration panel is a WebviewPanel but does not use the AI state machine.
+            RPCLayer._messenger.registerWebviewPanel(webViewPanel as WebviewPanel);
         } else {
             RPCLayer._messenger.registerWebviewView(webViewPanel as WebviewView);
             AIStateMachine.service().onTransition((state) => {
@@ -182,6 +186,10 @@ function isWebviewPanel(webview: WebviewPanel | WebviewView): boolean {
     return title === VisualizerWebview.webviewTitle;
 }
 
+function isMigrationPanel(webview: WebviewPanel | WebviewView): boolean {
+    return 'reveal' in webview && webview.title === "Migration Assistant";
+}
+
 export function notifyCurrentWebview() {
     RPCLayer._messenger.sendNotification(projectContentUpdated, { type: 'webview', webviewType: VisualizerWebview.viewType }, true);
 }
@@ -206,8 +214,8 @@ export function notifyApprovalOverlayState(state: ApprovalOverlayState) {
     RPCLayer._messenger.sendNotification(approvalOverlayState, { type: 'webview', webviewType: AiPanelWebview.viewType }, state);
 }
 
-export function notifyOnIdentifierUpdated(response: ProjectStructureArtifactResponse[]) {
-    RPCLayer._messenger.sendNotification(onIdentifierUpdated, { type: 'webview', webviewType: VisualizerWebview.viewType }, response);
+export function notifyOnIdentifierUpdated(artifacts: ProjectStructureArtifactResponse[]) {
+    RPCLayer._messenger.sendNotification(onIdentifierUpdated, { type: 'webview', webviewType: VisualizerWebview.viewType }, artifacts);
 }
 
 export function sendAgentChangedNotification(agentName: string) {

workspaces/ballerina/ballerina-extension/src/extension.ts

@@ -36,6 +36,7 @@ import { activate as activateLibraryBrowser } from './features/library-browser';
 import { activate as activateBIFeatures } from './features/bi';
 import { activate as activateERDiagram } from './views/persist-layer-diagram';
 import { activateAiPanel } from './views/ai-panel';
+import { activateMigrationPanel } from './views/migration-panel';
 import { debug, handleResolveMissingDependencies, log } from './utils';
 import { activateUriHandlers } from './utils/uri-handlers';
 import { StateMachine } from './stateMachine';
@@ -50,6 +51,7 @@ import { activate as activateNPFeatures } from './features/natural-programming/a
 import { activateAgentChatPanel } from './views/agent-chat/activate';
 import { activateTracing } from './features/tracing';
 import { activateICP } from './features/icp';
+import { onWizardChatNotify, setWizardProjectRoot, runWizardMigrationEnhancement, abortMigrationAgent, openMigratedProject } from './features/ai/migration/orchestrator';
 
 let langClient: ExtendedLangClient;
 export let isPluginStartup = true;
@@ -135,7 +137,14 @@ export async function activate(context: ExtensionContext) {
         ballerinaExtInstance: extension.ballerinaExtInstance, 
         projectPath: StateMachine.context().projectPath,
         VisualizerWebview,
-        BallerinaExtensionState
+        BallerinaExtensionState,
+        migration: {
+            setWizardProjectRoot,
+            wizardEnhancementReady: runWizardMigrationEnhancement,
+            abortAgent: abortMigrationAgent,
+            openMigratedProject,
+            onChatNotify: onWizardChatNotify,
+        },
     };
 }
 
@@ -202,6 +211,9 @@ export async function activateBallerina(): Promise<BallerinaExtension> {
         //activate ai panel
         activateAiPanel(ballerinaExtInstance);
 
+        // Activate migration enhancement panel
+        activateMigrationPanel(ballerinaExtInstance);
+
         // Activate AI features
         activateAIFeatures(ballerinaExtInstance);
 

workspaces/ballerina/ballerina-extension/src/features/ai/agent/AgentExecutor.ts

@@ -157,6 +157,9 @@ async function determineAffectedPackages(
  * - Plan approval workflow (via ApprovalManager in TaskWrite tool)
  */
 export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
+    /** Tracks in-flight tool-call start times keyed by toolCallId for duration logging. */
+    private readonly _pendingToolCalls = new Map<string, number>();
+
     constructor(config: AICommandConfig<GenerateAgentCodeRequest>) {
         super(config);
     }
@@ -253,6 +256,11 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
                 },
             ];
 
+            // Resolve model and limits
+            const llmModel = this.config.model ?? await getAnthropicClient(ANTHROPIC_SONNET_4);
+            const maxSteps = this.config.agentLimits?.maxSteps ?? 50;
+            const maxOutputTokens = this.config.agentLimits?.maxOutputTokens ?? 8192;
+
             // Create tools
             const tools = createToolRegistry({
                 eventHandler: this.config.eventHandler,
@@ -264,6 +272,7 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
                 projectRootPath: this.config.executionContext.workspacePath || this.config.executionContext.projectPath || '',
                 generationId: this.config.generationId,
                 threadId: 'default',
+                migrationSourcePath: this.config.toolOptions?.migrationSourcePath,
                 runningServices: runningServicesManager,
                 webSearchEnabled: params.webSearchEnabled ?? false,
                 ctx: this.config.executionContext,
@@ -455,6 +464,14 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                         }
                     );
 
+                    // Notify caller with partial messages (wizard migration history persistence)
+                    if (this.config.onMessagesAvailable) {
+                        this.config.onMessagesAvailable(
+                            [{ role: "user", content: streamContext.userMessageContent }, ...partialLLMMessages],
+                            'aborted'
+                        );
+                    }
+
                     // Note: Abort event is sent by base class handleExecutionError()
                 }
 
@@ -487,9 +504,17 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                 throw error;
             }
 
+            console.error("[AgentExecutor] Non-abort error in execute():", error);
+
+            // Abort the controller so that any in-flight tool executions
+            // managed by the AI SDK are cancelled and stop emitting events.
+            if (!this.config.abortController.signal.aborted) {
+                this.config.abortController.abort();
+            }
+
             this.config.eventHandler({
                 type: "error",
-                content: "An error occurred during agent execution. Please check the logs for details."
+                content: getErrorMessage(error)
             });
 
             // For other errors, return result with error
@@ -535,8 +560,31 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                 await this.handleStreamFinish(context);
                 break;
 
+            case "tool-call":
+                if (this.config.debugLogger) {
+                    this._pendingToolCalls.set((part as any).toolCallId, Date.now());
+                }
+                break;
+
+            case "tool-result":
+                if (this.config.debugLogger) {
+                    const startMs = this._pendingToolCalls.get((part as any).toolCallId);
+                    if (startMs !== undefined) {
+                        const durationMs = Date.now() - startMs;
+                        const rawResult = (part as any).result;
+                        const raw = typeof rawResult === "string"
+                            ? rawResult
+                            : rawResult === undefined
+                                ? "<no result>"
+                                : JSON.stringify(rawResult) ?? "<no result>";
+                        this.config.debugLogger.logToolCall((part as any).toolName, durationMs, raw);
+                        this._pendingToolCalls.delete((part as any).toolCallId);
+                    }
+                }
+                break;
+
             default:
-                // Tool calls/results handled automatically by SDK
+                // All other stream part types (step-finish, etc.) are handled by the SDK.
                 break;
         }
     }
@@ -682,27 +730,36 @@ Generation stopped by user. The last in-progress task was not saved. Files have
         // Update chat state storage
         await this.updateChatState(context, assistantMessages, tempProjectPath);
 
-        // Integrate generated code into workspace immediately so user sees changes during review
-        const workspaceId = context.ctx.workspacePath || context.ctx.projectPath;
-        const pendingReview = chatStateStorage.getPendingReviewGeneration(workspaceId, 'default');
-        if (pendingReview && pendingReview.reviewState.modifiedFiles.length > 0) {
-            const integrationCtx = { ...context.ctx };
-            // In workspace mode, resolve project path from modified files if not set
-            if (!integrationCtx.projectPath && integrationCtx.workspacePath && pendingReview.reviewState.modifiedFiles.length > 0) {
-                const firstBalFile = pendingReview.reviewState.modifiedFiles.find(f => f.endsWith('.bal'));
-                if (firstBalFile) {
-                    const packageName = firstBalFile.split('/')[0];
-                    if (packageName) {
-                        integrationCtx.projectPath = path.join(integrationCtx.workspacePath, packageName);
-                        StateMachine.context().projectPath = integrationCtx.projectPath;
+        // Integrate generated code into workspace immediately so user sees changes during review.
+        // Skip only when the temp path IS the real workspace directory (migration wizard in-place
+        // editing). A review-continuation run also sets existingTempPath but to a temp dir, so we
+        // compare resolved paths rather than just checking existingTempPath presence.
+        const workspaceRoot = context.ctx.workspacePath || context.ctx.projectPath;
+        const isInPlaceEditing =
+            !!workspaceRoot &&
+            path.resolve(tempProjectPath) === path.resolve(workspaceRoot);
+        if (!isInPlaceEditing) {
+            const workspaceId = workspaceRoot;
+            const pendingReview = chatStateStorage.getPendingReviewGeneration(workspaceId, 'default');
+            if (pendingReview && pendingReview.reviewState.modifiedFiles.length > 0) {
+                const integrationCtx = { ...context.ctx };
+                // In workspace mode, resolve project path from modified files if not set
+                if (!integrationCtx.projectPath && integrationCtx.workspacePath && pendingReview.reviewState.modifiedFiles.length > 0) {
+                    const firstBalFile = pendingReview.reviewState.modifiedFiles.find(f => f.endsWith('.bal'));
+                    if (firstBalFile) {
+                        const packageName = firstBalFile.split('/')[0];
+                        if (packageName) {
+                            integrationCtx.projectPath = path.join(integrationCtx.workspacePath, packageName);
+                            StateMachine.context().projectPath = integrationCtx.projectPath;
+                        }
                     }
                 }
+                await integrateCodeToWorkspace(
+                    pendingReview.reviewState.tempProjectPath!,
+                    new Set(pendingReview.reviewState.modifiedFiles),
+                    integrationCtx
+                );
             }
-            await integrateCodeToWorkspace(
-                pendingReview.reviewState.tempProjectPath!,
-                new Set(pendingReview.reviewState.modifiedFiles),
-                integrationCtx
-            );
         }
 
         // Emit UI events

workspaces/ballerina/ballerina-extension/src/features/ai/agent/tool-registry.ts

@@ -42,6 +42,12 @@ import { createConnectorGeneratorTool, CONNECTOR_GENERATOR_TOOL } from './tools/
 import { LIBRARY_SEARCH_TOOL, getLibrarySearchTool } from './tools/library-search';
 import { createConfigCollectorTool, CONFIG_COLLECTOR_TOOL } from './tools/config-collector';
 import { createTestRunnerTool, TEST_RUNNER_TOOL_NAME } from './tools/test-runner';
+import {
+    createMigrationSourceListTool,
+    createMigrationSourceReadTool,
+    MIGRATION_SOURCE_LIST_TOOL,
+    MIGRATION_SOURCE_READ_TOOL,
+} from './tools/migration-source-reader';
 import { createBallerinaRunTool, BALLERINA_RUN_TOOL_NAME } from './tools/ballerina-run';
 import { createBallerinaGetLogsTool, BALLERINA_GET_LOGS_TOOL_NAME } from './tools/ballerina-get-logs';
 import { createBallerinaStopTool, BALLERINA_STOP_TOOL_NAME } from './tools/ballerina-stop';
@@ -60,13 +66,15 @@ export interface ToolRegistryOptions {
     projectRootPath: string;
     generationId: string;
     threadId?: string;
+    /** Absolute path to the original migration source project (Mule, Tibco, etc.). */
+    migrationSourcePath?: string;
     runningServices: RunningServicesManager;
     webSearchEnabled: boolean;
     ctx: ExecutionContext;
 }
 
 export function createToolRegistry(opts: ToolRegistryOptions) {
-    const { eventHandler, tempProjectPath, modifiedFiles, allModifiedFiles, projects, generationType, projectRootPath, generationId, threadId, webSearchEnabled, ctx } = opts;
+    const { eventHandler, tempProjectPath, modifiedFiles, allModifiedFiles, projects, generationType, projectRootPath, generationId, threadId, migrationSourcePath, webSearchEnabled, ctx } = opts;
     return {
         [TASK_WRITE_TOOL_NAME]: createTaskWriteTool(
             eventHandler,
@@ -114,6 +122,11 @@ export function createToolRegistry(opts: ToolRegistryOptions) {
         ),
         [DIAGNOSTICS_TOOL_NAME]: createDiagnosticsTool(tempProjectPath, eventHandler),
         [TEST_RUNNER_TOOL_NAME]: createTestRunnerTool(tempProjectPath, eventHandler, modifiedFiles, allModifiedFiles, ctx),
+        // Migration source tools — registered only when a source project path is available
+        ...(migrationSourcePath ? {
+            [MIGRATION_SOURCE_LIST_TOOL]: createMigrationSourceListTool(eventHandler, migrationSourcePath),
+            [MIGRATION_SOURCE_READ_TOOL]: createMigrationSourceReadTool(eventHandler, migrationSourcePath),
+        } : {}),
         [HURL_TOOL_NAME]: createHurlTool(eventHandler),
         [BALLERINA_RUN_TOOL_NAME]: createBallerinaRunTool(tempProjectPath, opts.runningServices, eventHandler, modifiedFiles, allModifiedFiles, ctx),
         [BALLERINA_GET_LOGS_TOOL_NAME]: createBallerinaGetLogsTool(opts.runningServices, eventHandler),