Add diagnostic validation and fallback integration

workspaces/ballerina/ballerina-core/package.json

@@ -27,7 +27,8 @@
     "tree-kill": "^1.2.2",
     "vscode-uri": "^3.0.8",
     "@types/mousetrap": "~1.6.11",
-    "@types/ws": "^8.2.1"
+    "@types/ws": "^8.2.1",
+    "ai": "^5.0.101"
   },
   "devDependencies": {
     "@types/node": "^22.15.21",

workspaces/ballerina/ballerina-core/src/state-machine-types.ts

@@ -18,6 +18,7 @@
 
 import { NotificationType, RequestType } from "vscode-messenger-common";
 import { NodePosition, STNode } from "@wso2/syntax-tree";
+import { ModelMessage } from "ai";
 import { Command } from "./interfaces/ai-panel";
 import { LinePosition } from "./interfaces/common";
 import { Type } from "./interfaces/extended-lang-client";
@@ -434,7 +435,7 @@ export interface ChatMessage {
     id: string;
     content: string;
     uiResponse: string;
-    modelMessages: any[];
+    modelMessages: ModelMessage[];
     timestamp: number;
     checkpointId?: string;
 }

workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts

@@ -267,6 +267,14 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                     await integrateCodeToWorkspace(tempProjectPath, modifiedFilesSet);
                 }
 
+                // Fallback integration: integrate any remaining modified files that weren't integrated via TaskWrite in plan mode
+                if (isPlanModeEnabled && modifiedFiles.length > 0) {
+                    const modifiedFilesSet = new Set(modifiedFiles);
+                    await integrateCodeToWorkspace(tempProjectPath, modifiedFilesSet);
+                    console.log(`[Design] Successfully integrated files on stream completion`);
+                    modifiedFiles.length = 0;
+                }
+
                 updateAndSaveChat(messageId, userMessageContent, assistantMessages, eventHandler);
                 eventHandler({ type: "stop", command: Command.Design });
                 AIChatStateMachine.sendEvent({

workspaces/ballerina/ballerina-extension/src/features/ai/service/libs/task_write_tool.ts

@@ -17,9 +17,11 @@
 import { tool } from 'ai';
 import { z } from 'zod';
 import { CopilotEventHandler } from '../event';
-import { Task, TaskStatus, TaskTypes, Plan, AIChatMachineEventType, SourceFiles } from '@wso2/ballerina-core';
+import { Task, TaskStatus, TaskTypes, Plan, AIChatMachineEventType, SourceFiles, AIChatMachineContext } from '@wso2/ballerina-core';
 import { AIChatStateMachine } from '../../../../views/ai-panel/aiChatMachine';
 import { integrateCodeToWorkspace } from '../design/utils';
+import { checkCompilationErrors } from './diagnostics_utils';
+import { DIAGNOSTICS_TOOL_NAME } from './diagnostics_tool';
 
 export const TASK_WRITE_TOOL_NAME = "TaskWrite";
 
@@ -134,11 +136,11 @@ Rules:
 
                 const taskCategories = categorizeTasks(allTasks);
 
-                // TODO: Fix issue where agent updates to existing plan trigger new approval
-                // Problem: When agent continues chat with plan updates, currentPlan state is empty
-                // causing it to be identified as new plan and triggering approval unnecessarily.
-                // Need to preserve plan state across chat continuations or use chat history to
-                // detect if this is a continuation of an existing conversation with a plan.
+                // TODO: Add tests for plan modification detection in the middle of execution
+                // Fixed: Plan state is now preserved in the state machine across chat continuations,
+                // preventing unnecessary approval requests when agent continues with existing plan.
+                // Still need comprehensive tests for: mid-execution plan modifications, task reordering,
+                // task additions/removals, and edge cases where plan changes should trigger re-approval.
                 const isNewPlan = !existingPlan || existingPlan.tasks.length === 0;
                 const isPlanRemodification = existingPlan && (
                     allTasks.length !== existingPlan.tasks.length ||
@@ -304,18 +306,32 @@ async function handlePlanApproval(
 async function handleTaskCompletion(
     allTasks: Task[],
     newlyCompletedTasks: Task[],
-    currentContext: any,
+    currentContext: AIChatMachineContext,
     eventHandler: CopilotEventHandler,
     tempProjectPath?: string,
     modifiedFiles?: string[]
 ): Promise<{ approved: boolean; comment?: string; approvedTaskDescription: string }> {
     const lastCompletedTask = newlyCompletedTasks[newlyCompletedTasks.length - 1];
     console.log(`[TaskWrite Tool] Detected ${newlyCompletedTasks.length} newly completed task(s)`);
 
+    const diagnosticResult = await checkCompilationErrors(tempProjectPath!);
+
+    if (diagnosticResult.diagnostics.length > 0) {
+        const errorCount = diagnosticResult.diagnostics.length;
+        console.error(`[TaskWrite Tool] Found ${errorCount} compilation error(s), blocking task completion`);
+
+        return {
+            approved: false,
+            comment: `Cannot complete task: ${errorCount} compilation error(s) detected. Use the ${DIAGNOSTICS_TOOL_NAME} tool to check the errors and fix them before marking the task as completed.`,
+            approvedTaskDescription: lastCompletedTask.description
+        };
+    }
+
     if (tempProjectPath && modifiedFiles) {
         const modifiedFilesSet = new Set(modifiedFiles);
         console.log(`[TaskWrite Tool] Integrating ${modifiedFilesSet.size} modified file(s)`);
         await integrateCodeToWorkspace(tempProjectPath, modifiedFilesSet);
+        modifiedFiles.length = 0;
     }
 
     AIChatStateMachine.sendEvent({

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx

@@ -121,6 +121,7 @@ const convertToUIMessages = (messages: UIChatHistoryMessage[]) => {
 };
 
 // Helper function to load chat history from localStorage
+// TODO: Need to be removed and move chat history to statemachine fully
 const loadFromLocalStorage = (projectUuid: string, setMessages: React.Dispatch<React.SetStateAction<any[]>>, chatArray: ChatEntry[]) => {
     const localStorageFile = `chatArray-AIGenerationChat-${projectUuid}`;
     const storedChatArray = localStorage.getItem(localStorageFile);
@@ -315,6 +316,7 @@ const AIChat: React.FC = () => {
             setMessages(uiMessages);
 
             chatArray = chatArray.slice(0, updatedMessages.length);
+            // TODO: Need to be removed and move chat history to statemachine fully
             localStorage.setItem(`chatArray-AIGenerationChat-${projectUuid}`, JSON.stringify(chatArray));
 
             setIsLoading(false);
@@ -1367,6 +1369,7 @@ const AIChat: React.FC = () => {
                 filepath: chatLocation,
             });
             integratedChatIndex = previouslyIntegratedChatIndex;
+            // TODO: Need to be removed and move chat history to statemachine fully
             localStorage.setItem(
                 `chatArray-AIGenerationChat-${projectUuid}-developer-index`,
                 JSON.stringify({ integratedChatIndex, previouslyIntegratedChatIndex })