Add proper TypeScript types and AI SDK dependency

- Add AIChatMachineContext type to task_write_tool.ts (replaces any)
- Add ModelMessage type to state-machine-types.ts (replaces any[])
- Add AI SDK v5.0.101 dependency to ballerina-core for ModelMessage type
- Add TODO comments for localStorage migration to state machine

Author: RNViththagan

common/config/rush/pnpm-lock.yaml

@@ -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/package.json

@@ -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-core/src/state-machine-types.ts

@@ -17,7 +17,7 @@
 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';
@@ -136,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 ||
@@ -306,7 +306,7 @@ async function handlePlanApproval(
 async function handleTaskCompletion(
     allTasks: Task[],
     newlyCompletedTasks: Task[],
-    currentContext: any,
+    currentContext: AIChatMachineContext,
     eventHandler: CopilotEventHandler,
     tempProjectPath?: string,
     modifiedFiles?: string[]

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

@@ -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 })