diff --git a/common/config/rush/pnpm-lock.yaml b/common/config/rush/pnpm-lock.yaml index 46dad21a04..1a17aeb806 100644 --- a/common/config/rush/pnpm-lock.yaml +++ b/common/config/rush/pnpm-lock.yaml @@ -385,6 +385,9 @@ importers: '@wso2/syntax-tree': specifier: workspace:* version: link:../syntax-tree + ai: + specifier: ^5.0.101 + version: 5.0.101(zod@4.1.11) handlebars: specifier: ~4.7.8 version: 4.7.8 diff --git a/workspaces/ballerina/ballerina-core/package.json b/workspaces/ballerina/ballerina-core/package.json index bdec860366..3c9a92ecca 100644 --- a/workspaces/ballerina/ballerina-core/package.json +++ b/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", diff --git a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts index ea06aa6dc5..3c06ceddcf 100644 --- a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts +++ b/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; } diff --git a/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts b/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts index 16f2ad3611..1bcadf8291 100644 --- a/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts +++ b/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({ diff --git a/workspaces/ballerina/ballerina-extension/src/features/ai/service/libs/task_write_tool.ts b/workspaces/ballerina/ballerina-extension/src/features/ai/service/libs/task_write_tool.ts index 66c9d6d8c3..62aeaadea2 100644 --- a/workspaces/ballerina/ballerina-extension/src/features/ai/service/libs/task_write_tool.ts +++ b/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,7 +306,7 @@ async function handlePlanApproval( async function handleTaskCompletion( allTasks: Task[], newlyCompletedTasks: Task[], - currentContext: any, + currentContext: AIChatMachineContext, eventHandler: CopilotEventHandler, tempProjectPath?: string, modifiedFiles?: string[] @@ -312,10 +314,24 @@ async function handleTaskCompletion( 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({ diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx index 10444aea5a..f8071b7512 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx +++ b/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>, 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 })