From c140bae41a0db09ba840c1d74000bc9df0e15e6c Mon Sep 17 00:00:00 2001 From: Anjana Supun Date: Fri, 17 Oct 2025 13:39:01 +0530 Subject: [PATCH 001/619] Add /design command --- .../ballerina-core/src/interfaces/ai-panel.ts | 1 + .../src/rpc-types/ai-panel/index.ts | 3 +- .../src/rpc-types/ai-panel/rpc-type.ts | 3 +- .../src/features/ai/service/design/design.ts | 80 +++++++++++++++++++ .../src/rpc-managers/ai-panel/rpc-handler.ts | 4 +- .../src/rpc-managers/ai-panel/rpc-manager.ts | 13 ++- .../src/rpc-clients/ai-panel/rpc-client.ts | 7 +- .../data/commandTemplates.const.ts | 13 +++ .../data/placeholderTags.const.ts | 3 + .../views/AIPanel/components/AIChat/index.tsx | 19 +++++ 10 files changed, 139 insertions(+), 7 deletions(-) create mode 100644 workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts diff --git a/workspaces/ballerina/ballerina-core/src/interfaces/ai-panel.ts b/workspaces/ballerina/ballerina-core/src/interfaces/ai-panel.ts index 54415539d0a..d8ac17a2e8e 100644 --- a/workspaces/ballerina/ballerina-core/src/interfaces/ai-panel.ts +++ b/workspaces/ballerina/ballerina-core/src/interfaces/ai-panel.ts @@ -25,6 +25,7 @@ export enum Command { Ask = '/ask', NaturalProgramming = '/natural-programming (experimental)', OpenAPI = '/openapi', + Design = '/design', Doc = '/doc' } diff --git a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts index f0e1ad75686..4a3a3b3cece 100644 --- a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts +++ b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts @@ -79,6 +79,7 @@ export interface AIPanelAPI { getRelevantLibrariesAndFunctions: (params: RelevantLibrariesAndFunctionsRequest) => Promise; generateOpenAPI: (params: GenerateOpenAPIRequest) => void; generateCode: (params: GenerateCodeRequest) => void; + generateDesign: (params: GenerateCodeRequest) => Promise; repairGeneratedCode: (params: RepairParams) => void; generateTestPlan: (params: TestPlanGenerationRequest) => void; generateFunctionTests: (params: TestGeneratorIntermediaryState) => void; @@ -87,6 +88,6 @@ export interface AIPanelAPI { // ================================== // Doc Generation Related Functions // ================================== - getGeneratedDocumentation: (params: DocGenerationRequest) => Promise; + getGeneratedDocumentation: (params: DocGenerationRequest) => Promise; addFilesToProject: (params: AddFilesToProjectRequest) => Promise; } diff --git a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts index 6b6140ce5fe..6d9e1b5bb68 100644 --- a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts +++ b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts @@ -78,10 +78,11 @@ export const submitFeedback: RequestType = { met export const getRelevantLibrariesAndFunctions: RequestType = { method: `${_preFix}/getRelevantLibrariesAndFunctions` }; export const generateOpenAPI: NotificationType = { method: `${_preFix}/generateOpenAPI` }; export const generateCode: NotificationType = { method: `${_preFix}/generateCode` }; +export const generateDesign: RequestType = { method: `${_preFix}/generateDesign` }; export const repairGeneratedCode: NotificationType = { method: `${_preFix}/repairGeneratedCode` }; export const generateTestPlan: NotificationType = { method: `${_preFix}/generateTestPlan` }; export const generateFunctionTests: NotificationType = { method: `${_preFix}/generateFunctionTests` }; export const generateHealthcareCode: NotificationType = { method: `${_preFix}/generateHealthcareCode` }; export const abortAIGeneration: NotificationType = { method: `${_preFix}/abortAIGeneration` }; -export const getGeneratedDocumentation: NotificationType = { method: `${_preFix}/getGeneratedDocumentation` }; +export const getGeneratedDocumentation: RequestType = { method: `${_preFix}/getGeneratedDocumentation` }; export const addFilesToProject: RequestType = { method: `${_preFix}/addFilesToProject` }; 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 new file mode 100644 index 00000000000..1df36a487f9 --- /dev/null +++ b/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts @@ -0,0 +1,80 @@ +// Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com/) All Rights Reserved. + +// WSO2 LLC. licenses this file to you under the Apache License, +// Version 2.0 (the "License"); you may not use this file except +// in compliance with the License. +// You may obtain a copy of the License at + +// http://www.apache.org/licenses/LICENSE-2.0 + +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +import { Command, GenerateCodeRequest } from "@wso2/ballerina-core"; +import { streamText } from "ai"; +import { getAnthropicClient, ANTHROPIC_HAIKU, getProviderCacheControl } from "../connection"; +import { getErrorMessage, populateHistory } from "../utils"; +import { CopilotEventHandler, createWebviewEventHandler } from "../event"; +import { AIPanelAbortController } from "../../../../rpc-managers/ai-panel/utils"; + +export async function generateDesignCore(params: GenerateCodeRequest, eventHandler: CopilotEventHandler): Promise { + // Populate chat history and add user message + const historyMessages = populateHistory(params.chatHistory); + const cacheOptions = await getProviderCacheControl(); + const { fullStream } = streamText({ + model: await getAnthropicClient(ANTHROPIC_HAIKU), + maxOutputTokens: 8192, + temperature: 0, + messages: [ + { + role: "system", + content: "Your task is to generate a highlevel design for a Ballerina integration based on the user's requirements. Provide a structured design outline including components, data flow, and interactions.", + providerOptions: cacheOptions + }, + ...historyMessages, + { + role: "user", + content: params.usecase + }, + ], + abortSignal: AIPanelAbortController.getInstance().signal, + }); + + eventHandler({ type: "start" }); + + for await (const part of fullStream) { + switch (part.type) { + case "text-delta": { + const textPart = part.text; + eventHandler({ type: "content_block", content: textPart }); + break; + } + case "error": { + const error = part.error; + console.error("Error during design generation:", error); + eventHandler({ type: "error", content: getErrorMessage(error) }); + break; + } + case "finish": { + const finishReason = part.finishReason; + eventHandler({ type: "stop", command: Command.Design }); + break; + } + } + } +} + +// Main public function that uses the default event handler +export async function generatDesign(params: GenerateCodeRequest): Promise { + const eventHandler = createWebviewEventHandler(Command.Design); + try { + await generateDesignCore(params, eventHandler); + } catch (error) { + console.error("Error during design generation:", error); + eventHandler({ type: "error", content: getErrorMessage(error) }); + } +} diff --git a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts index 8404ed51eb1..1aca64b329b 100644 --- a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts +++ b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts @@ -46,6 +46,7 @@ import { generateCode, GenerateCodeRequest, generateDataMapperModel, + generateDesign, generateFunctionTests, generateHealthcareCode, generateMappings, @@ -169,11 +170,12 @@ export function registerAiPanelRpcHandlers(messenger: Messenger) { messenger.onRequest(getRelevantLibrariesAndFunctions, (args: RelevantLibrariesAndFunctionsRequest) => rpcManger.getRelevantLibrariesAndFunctions(args)); messenger.onNotification(generateOpenAPI, (args: GenerateOpenAPIRequest) => rpcManger.generateOpenAPI(args)); messenger.onNotification(generateCode, (args: GenerateCodeRequest) => rpcManger.generateCode(args)); + messenger.onRequest(generateDesign, (args: GenerateCodeRequest) => rpcManger.generateDesign(args)); messenger.onNotification(repairGeneratedCode, (args: RepairParams) => rpcManger.repairGeneratedCode(args)); messenger.onNotification(generateTestPlan, (args: TestPlanGenerationRequest) => rpcManger.generateTestPlan(args)); messenger.onNotification(generateFunctionTests, (args: TestGeneratorIntermediaryState) => rpcManger.generateFunctionTests(args)); messenger.onNotification(generateHealthcareCode, (args: GenerateCodeRequest) => rpcManger.generateHealthcareCode(args)); messenger.onNotification(abortAIGeneration, () => rpcManger.abortAIGeneration()); - messenger.onNotification(getGeneratedDocumentation, (args: DocGenerationRequest) => rpcManger.getGeneratedDocumentation(args)); + messenger.onRequest(getGeneratedDocumentation, (args: DocGenerationRequest) => rpcManger.getGeneratedDocumentation(args)); messenger.onRequest(addFilesToProject, (args: AddFilesToProjectRequest) => rpcManger.addFilesToProject(args)); } diff --git a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts index f74490e5f92..7d460d27328 100644 --- a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts +++ b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts @@ -74,7 +74,7 @@ import { TestGenerationResponse, TestGeneratorIntermediaryState, TestPlanGenerationRequest, - TextEdit + TextEdit, } from "@wso2/ballerina-core"; import * as crypto from 'crypto'; import * as fs from 'fs'; @@ -87,6 +87,7 @@ import { FunctionDefinition, ModulePart, STKindChecker, STNode } from "@wso2/syn import { isNumber } from "lodash"; import { URI } from "vscode-uri"; import { NOT_SUPPORTED } from "../../../src/core/extended-language-client"; +import { CLOSE_AI_PANEL_COMMAND, OPEN_AI_PANEL_COMMAND } from "../../../src/features/ai/constants"; import { fetchWithAuth } from "../../../src/features/ai/service/connection"; import { generateOpenAPISpec } from "../../../src/features/ai/service/openapi/openapi"; import { AIStateMachine } from "../../../src/views/ai-panel/aiMachine"; @@ -118,7 +119,7 @@ import { import { attemptRepairProject, checkProjectDiagnostics } from "./repair-utils"; import { AIPanelAbortController, addToIntegration, cleanDiagnosticMessages, handleStop, isErrorCode, processMappings, requirementsSpecification, searchDocumentation } from "./utils"; import { fetchData } from "./utils/fetch-data-utils"; -import { CLOSE_AI_PANEL_COMMAND, OPEN_AI_PANEL_COMMAND } from "../../../src/features/ai/constants"; +import { generatDesign, generateDesignCore } from "../../../src/features/ai/service/design/design"; export class AiPanelRpcManager implements AIPanelAPI { @@ -1014,8 +1015,9 @@ export class AiPanelRpcManager implements AIPanelAPI { } } - async getGeneratedDocumentation(params: DocGenerationRequest): Promise { + async getGeneratedDocumentation(params: DocGenerationRequest): Promise { await generateDocumentationForService(params); + return true; } async addFilesToProject(params: AddFilesToProjectRequest): Promise { @@ -1039,6 +1041,11 @@ export class AiPanelRpcManager implements AIPanelAPI { return false; //silently fail for timeout issues. } } + + async generateDesign(params: GenerateCodeRequest): Promise { + await generatDesign(params); + return true; + } } function getModifiedAssistantResponse(originalAssistantResponse: string, tempDir: string, project: ProjectSource): string { diff --git a/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts b/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts index b70e51a425a..bed89f599c1 100644 --- a/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts +++ b/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts @@ -76,6 +76,7 @@ import { fetchData, generateCode, generateDataMapperModel, + generateDesign, generateFunctionTests, generateHealthcareCode, generateMappings, @@ -353,6 +354,10 @@ export class AiPanelRpcClient implements AIPanelAPI { return this._messenger.sendNotification(generateCode, HOST_EXTENSION, params); } + generateDesign(params: GenerateCodeRequest): Promise { + return this._messenger.sendRequest(generateDesign, HOST_EXTENSION, params); + } + repairGeneratedCode(params: RepairParams): void { return this._messenger.sendNotification(repairGeneratedCode, HOST_EXTENSION, params); } @@ -373,7 +378,7 @@ export class AiPanelRpcClient implements AIPanelAPI { return this._messenger.sendNotification(abortAIGeneration, HOST_EXTENSION); } - getGeneratedDocumentation(params: DocGenerationRequest): Promise { + getGeneratedDocumentation(params: DocGenerationRequest): Promise { return this._messenger.sendRequest(getGeneratedDocumentation, HOST_EXTENSION, params); } diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts index f6ba0905cbe..1f21f451ef7 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts @@ -139,6 +139,13 @@ export const commandTemplates = { placeholders: [], }, ], + [Command.Design]: [ + { + id: TemplateId.Wildcard, + text: '', + placeholders: [], + }, + ], [Command.Doc]: [ { id: TemplateId.GenerateUserDoc, @@ -194,6 +201,12 @@ export const suggestedCommandTemplates: AIPanelPrompt[] = [ templateId: TemplateId.Wildcard, text: 'write a hello world http service', }, + { + type: 'command-template', + command: Command.Design, + templateId: TemplateId.Wildcard, + text: 'design an API for a task management system', + }, { type: 'command-template', command: Command.Ask, diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/placeholderTags.const.ts b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/placeholderTags.const.ts index 8b7dad97cfa..15b599da0eb 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/placeholderTags.const.ts +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/placeholderTags.const.ts @@ -72,6 +72,9 @@ export const placeholderTags: PlaceholderTagMap = { [Command.OpenAPI]: { 'wildcard': {}, }, + [Command.Design]: { + 'wildcard': {}, + }, [Command.Doc]: { 'generate-user-doc': { servicename: [], 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 6cf1eb320e5..98fc0b30f38 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 @@ -723,6 +723,14 @@ const AIChat: React.FC = () => { } break; } + case Command.Design: { + switch (parsedInput.templateId) { + case TemplateId.Wildcard: + await processDesignGeneration(parsedInput.text, inputText); + break; + } + break; + } case Command.Doc: { switch (parsedInput.templateId) { case "generate-user-doc": @@ -1669,6 +1677,17 @@ const AIChat: React.FC = () => { await rpcClient.getAiPanelRpcClient().generateOpenAPI(requestBody); } + async function processDesignGeneration(useCase: string, message: string) { + const requestBody: GenerateCodeRequest = { + usecase: useCase, + chatHistory: chatArray, + operationType: CodeGenerationType.CODE_GENERATION, + fileAttachmentContents: [], + }; + + await rpcClient.getAiPanelRpcClient().generateDesign(requestBody); + } + async function handleStop() { // Abort any ongoing requests // abortFetchWithAuth(); From 26c71bc743950d1f4b5cd13fc829ead6a8f288ab Mon Sep 17 00:00:00 2001 From: gigara Date: Mon, 20 Oct 2025 14:47:40 +0530 Subject: [PATCH 002/619] Add aiChat state machine --- .../ballerina-core/src/state-machine-types.ts | 92 +++ .../ballerina-extension/src/RPCLayer.ts | 7 +- .../ballerina-extension/src/stateMachine.ts | 2 + .../src/views/ai-panel/aiChatMachine.ts | 770 ++++++++++++++++++ .../src/BallerinaRpcClient.ts | 15 +- .../views/AIPanel/components/AIChat/index.tsx | 26 +- 6 files changed, 900 insertions(+), 12 deletions(-) create mode 100644 workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts diff --git a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts index f29c0387b09..66ab365a44d 100644 --- a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts +++ b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts @@ -324,6 +324,95 @@ export type AIMachineSendableEvent = : { type: K; payload: AIMachineEventMap[K] } }[keyof AIMachineEventMap]; +export type AIChatMachineStateValue = + | 'Idle' + | 'CreatingPlan' + | 'PlanReview' + | 'UpdatingPlan' + | 'ExecutingTask' + | 'AwaitingUserAction' + | 'FinetuningTask' + | 'QuestioningUser' + | 'Completed' + | 'Error'; + +export enum AIChatMachineEventType { + SUBMIT_PROMPT = 'SUBMIT_PROMPT', + PLAN_CREATED = 'PLAN_CREATED', + EDIT_TASK = 'EDIT_TASK', + UPDATE_PLAN_WITH_PROMPT = 'UPDATE_PLAN_WITH_PROMPT', + FINALIZE_PLAN = 'FINALIZE_PLAN', + TASK_COMPLETED = 'TASK_COMPLETED', + CONTINUE_TO_NEXT = 'CONTINUE_TO_NEXT', + FINETUNE_TASK = 'FINETUNE_TASK', + FINETUNE_COMPLETED = 'FINETUNE_COMPLETED', + ASK_QUESTION = 'ASK_QUESTION', + ANSWER_QUESTION = 'ANSWER_QUESTION', + ALL_TASKS_COMPLETED = 'ALL_TASKS_COMPLETED', + RESET = 'RESET', + RESTORE_STATE = 'RESTORE_STATE', + ERROR = 'ERROR', + RETRY = 'RETRY', +} + +export interface ChatMessage { + id: string; + role: 'user' | 'assistant' | 'system'; + content: string; + timestamp: number; +} + +export interface Task { + id: string; + description: string; + status: 'pending' | 'in-progress' | 'completed' | 'failed'; + result?: string; + error?: string; +} + +export interface Plan { + id: string; + tasks: Task[]; + createdAt: number; + updatedAt: number; +} + +export interface Question { + id: string; + question: string; + context?: string; + timestamp: number; +} + +export interface AIChatMachineContext { + initialPrompt?: string; + chatHistory: ChatMessage[]; + currentPlan?: Plan; + currentTaskIndex: number; + currentQuestion?: Question; + errorMessage?: string; + sessionId?: string; + projectId?: string; +} + +export type AIChatMachineSendableEvent = + | { type: AIChatMachineEventType.SUBMIT_PROMPT; payload: { prompt: string } } + | { type: AIChatMachineEventType.PLAN_CREATED; payload: { plan: Plan } } + | { type: AIChatMachineEventType.EDIT_TASK; payload: { taskId: string; description: string } } + | { type: AIChatMachineEventType.UPDATE_PLAN_WITH_PROMPT; payload: { prompt: string } } + | { type: AIChatMachineEventType.FINALIZE_PLAN } + | { type: AIChatMachineEventType.TASK_COMPLETED; payload: { result: string } } + | { type: AIChatMachineEventType.CONTINUE_TO_NEXT } + | { type: AIChatMachineEventType.FINETUNE_TASK; payload: { instructions: string } } + | { type: AIChatMachineEventType.FINETUNE_COMPLETED; payload: { result: string } } + | { type: AIChatMachineEventType.ASK_QUESTION; payload: { question: Question } } + | { type: AIChatMachineEventType.ANSWER_QUESTION; payload: { answer: string } } + | { type: AIChatMachineEventType.ALL_TASKS_COMPLETED } + | { type: AIChatMachineEventType.RESET } + | { type: AIChatMachineEventType.RESTORE_STATE; payload: { state: AIChatMachineContext } } + | { type: AIChatMachineEventType.ERROR; payload: { message: string } } + | { type: AIChatMachineEventType.RETRY }; + export enum LoginMethod { BI_INTEL = 'biIntel', ANTHROPIC_KEY = 'anthropic_key', @@ -380,3 +469,6 @@ export enum ColorThemeKind { export const aiStateChanged: NotificationType = { method: 'aiStateChanged' }; export const sendAIStateEvent: RequestType = { method: 'sendAIStateEvent' }; export const currentThemeChanged: NotificationType = { method: 'currentThemeChanged' }; + +export const aiChatStateChanged: NotificationType = { method: 'aiChatStateChanged' }; +export const sendAIChatStateEvent: RequestType = { method: 'sendAIChatStateEvent' }; \ No newline at end of file diff --git a/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts b/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts index 42dea359982..b074dda074d 100644 --- a/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts +++ b/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts @@ -19,7 +19,7 @@ import { WebviewView, WebviewPanel, window } from 'vscode'; import { Messenger } from 'vscode-messenger'; import { StateMachine } from './stateMachine'; -import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent } from '@wso2/ballerina-core'; +import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent, aiChatStateChanged, sendAIChatStateEvent, AIChatMachineEventType, AIChatMachineSendableEvent } from '@wso2/ballerina-core'; import { VisualizerWebview } from './views/visualizer/webview'; import { registerVisualizerRpcHandlers } from './rpc-managers/visualizer/rpc-handler'; import { registerLangClientRpcHandlers } from './rpc-managers/lang-client/rpc-handler'; @@ -45,6 +45,7 @@ import { extension } from './BalExtensionContext'; import { registerAgentChatRpcHandlers } from './rpc-managers/agent-chat/rpc-handler'; import { ArtifactsUpdated, ArtifactNotificationHandler } from './utils/project-artifacts-handler'; import { registerMigrateIntegrationRpcHandlers } from './rpc-managers/migrate-integration/rpc-handler'; +import { AIChatStateMachine } from './views/ai-panel/aiChatMachine'; export class RPCLayer { static _messenger: Messenger = new Messenger(); @@ -67,6 +68,9 @@ export class RPCLayer { AIStateMachine.service().onTransition((state) => { RPCLayer._messenger.sendNotification(aiStateChanged, { type: 'webview', webviewType: AiPanelWebview.viewType }, state.value); }); + AIChatStateMachine.service().onTransition((state) => { + RPCLayer._messenger.sendNotification(aiChatStateChanged, { type: 'webview', webviewType: AiPanelWebview.viewType }, state.value); + }); } } @@ -96,6 +100,7 @@ export class RPCLayer { // ----- AI Webview RPC Methods registerAiPanelRpcHandlers(RPCLayer._messenger); RPCLayer._messenger.onRequest(sendAIStateEvent, (event: AIMachineEventType | AIMachineSendableEvent) => AIStateMachine.sendEvent(event)); + RPCLayer._messenger.onRequest(sendAIChatStateEvent, (event: AIChatMachineEventType | AIChatMachineSendableEvent) => AIChatStateMachine.sendEvent(event)); // ----- Data Mapper Webview RPC Methods registerDataMapperRpcHandlers(RPCLayer._messenger); diff --git a/workspaces/ballerina/ballerina-extension/src/stateMachine.ts b/workspaces/ballerina/ballerina-extension/src/stateMachine.ts index dbc86d127c5..5058a1bce56 100644 --- a/workspaces/ballerina/ballerina-extension/src/stateMachine.ts +++ b/workspaces/ballerina/ballerina-extension/src/stateMachine.ts @@ -16,6 +16,7 @@ import { AIStateMachine } from './views/ai-panel/aiMachine'; import { StateMachinePopup } from './stateMachinePopup'; import { checkIsBallerina, checkIsBI, fetchScope, getOrgPackageName, UndoRedoManager } from './utils'; import { buildProjectArtifactsStructure } from './utils/project-artifacts'; +import { AIChatStateMachine } from './views/ai-panel/aiChatMachine'; interface MachineContext extends VisualizerLocation { langClient: ExtendedLangClient | null; @@ -307,6 +308,7 @@ const stateMachine = createMachine( const ls = await activateBallerina(); fetchAndCacheLibraryData(); AIStateMachine.initialize(); + AIChatStateMachine.initialize(); StateMachinePopup.initialize(); commands.executeCommand('setContext', 'BI.status', 'loadingDone'); if (!ls.biSupported) { diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts new file mode 100644 index 00000000000..9c93f4f07b8 --- /dev/null +++ b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts @@ -0,0 +1,770 @@ +/** + * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com) All Rights Reserved. + * + * WSO2 LLC. licenses this file to you under the Apache License, + * Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* eslint-disable @typescript-eslint/naming-convention */ +import { createMachine, assign, interpret } from 'xstate'; +import { extension } from '../../BalExtensionContext'; +import * as crypto from 'crypto'; +import { AIChatMachineContext, AIChatMachineEventType, AIChatMachineSendableEvent, AIChatMachineStateValue, ChatMessage, Plan, Task } from '@wso2/ballerina-core/lib/state-machine-types'; +import { workspace } from 'vscode'; +import { GenerateCodeRequest } from '@wso2/ballerina-core/lib/rpc-types/ai-panel/interfaces'; +import { generatDesign } from '../../features/ai/service/design/design'; + +const generateId = () => `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; + +const generateSessionId = () => `session-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; + +/** + * Generates a unique project identifier based on the workspace root path + * @returns A UUID string for the current project + */ +export const generateProjectId = (): string => { + const workspaceFolders = workspace.workspaceFolders; + + if (!workspaceFolders || workspaceFolders.length === 0) { + // Fallback for when no workspace is open + return 'default-project'; + } + + // Use the first workspace folder path to generate a consistent UUID + const workspacePath = workspaceFolders[0].uri.fsPath; + + // Create a hash of the workspace path for consistent project ID + const hash = crypto.createHash('sha256'); + hash.update(workspacePath); + const projectHash = hash.digest('hex').substring(0, 16); + + return `project-${projectHash}`; +}; + +const addChatMessage = ( + history: ChatMessage[], + role: 'user' | 'assistant' | 'system', + content: string +): ChatMessage[] => { + return [ + ...history, + { + id: generateId(), + role, + content, + timestamp: Date.now(), + }, + ]; +}; + +const chatMachine = createMachine({ + id: 'ballerina-ai-chat', + initial: 'Idle', + predictableActionArguments: true, + context: { + initialPrompt: undefined, + chatHistory: [], + currentPlan: undefined, + currentTaskIndex: -1, + currentQuestion: undefined, + errorMessage: undefined, + sessionId: undefined, + projectId: undefined, + }, + on: { + [AIChatMachineEventType.RESET]: { + target: 'Idle', + actions: [ + 'clearChatState', + assign({ + initialPrompt: (_ctx) => undefined, + chatHistory: (_ctx) => [], + currentPlan: (_ctx) => undefined, + currentTaskIndex: (_ctx) => -1, + currentQuestion: (_ctx) => undefined, + errorMessage: (_ctx) => undefined, + sessionId: (_ctx) => undefined, + // Keep projectId to maintain project context + }), + ], + }, + [AIChatMachineEventType.RESTORE_STATE]: { + target: 'PlanReview', + actions: assign({ + initialPrompt: (_ctx, event) => event.payload.state.initialPrompt, + chatHistory: (_ctx, event) => event.payload.state.chatHistory, + currentPlan: (_ctx, event) => event.payload.state.currentPlan, + currentTaskIndex: (_ctx, event) => event.payload.state.currentTaskIndex, + currentQuestion: (_ctx, event) => event.payload.state.currentQuestion, + errorMessage: (_ctx) => undefined, + sessionId: (_ctx, event) => event.payload.state.sessionId, + }), + }, + [AIChatMachineEventType.ERROR]: { + target: 'Error', + actions: assign({ + errorMessage: (_ctx, event) => event.payload.message, + }), + }, + }, + states: { + Idle: { + entry: assign({ + sessionId: (_ctx) => generateSessionId(), + projectId: (_ctx) => generateProjectId(), + }), + on: { + [AIChatMachineEventType.SUBMIT_PROMPT]: { + target: 'CreatingPlan', + actions: assign({ + initialPrompt: (_ctx, event) => event.payload.prompt, + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), + errorMessage: (_ctx) => undefined, + }), + }, + }, + }, + CreatingPlan: { + invoke: { + id: 'createPlan', + src: 'createPlan', + onDone: { + target: 'PlanReview', + actions: assign({ + currentPlan: (_ctx, event) => event.data.plan, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Plan created with ${event.data.plan.tasks.length} tasks` + ), + }), + }, + onError: { + target: 'Error', + actions: assign({ + errorMessage: (_ctx, event) => event.data?.message || 'Failed to create plan', + }), + }, + }, + }, + PlanReview: { + entry: 'saveChatState', + on: { + [AIChatMachineEventType.EDIT_TASK]: { + actions: assign({ + currentPlan: (ctx, event) => { + if (!ctx.currentPlan) { + return ctx.currentPlan; + } + return { + ...ctx.currentPlan, + tasks: ctx.currentPlan.tasks.map((task) => + task.id === event.payload.taskId + ? { ...task, description: event.payload.description } + : task + ), + updatedAt: Date.now(), + }; + }, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'user', + `Updated task: ${event.payload.description}` + ), + }), + }, + [AIChatMachineEventType.UPDATE_PLAN_WITH_PROMPT]: { + target: 'UpdatingPlan', + actions: assign({ + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), + }), + }, + [AIChatMachineEventType.FINALIZE_PLAN]: { + target: 'ExecutingTask', + actions: assign({ + currentTaskIndex: (_ctx) => 0, + chatHistory: (ctx) => + addChatMessage(ctx.chatHistory, 'system', 'Plan finalized. Starting execution...'), + }), + }, + }, + }, + UpdatingPlan: { + invoke: { + id: 'updatePlan', + src: 'updatePlan', + onDone: { + target: 'PlanReview', + actions: assign({ + currentPlan: (_ctx, event) => event.data.plan, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Plan updated with ${event.data.plan.tasks.length} tasks` + ), + }), + }, + onError: { + target: 'Error', + actions: assign({ + errorMessage: (_ctx, event) => event.data?.message || 'Failed to update plan', + }), + }, + }, + }, + ExecutingTask: { + entry: 'saveChatState', + invoke: { + id: 'executeTask', + src: 'executeTask', + onDone: { + target: 'AwaitingUserAction', + actions: assign({ + currentPlan: (ctx, event) => { + if (!ctx.currentPlan) { + return ctx.currentPlan; + } + return { + ...ctx.currentPlan, + tasks: ctx.currentPlan.tasks.map((task, index) => + index === ctx.currentTaskIndex + ? { ...task, status: 'completed', result: event.data.result } + : task + ), + updatedAt: Date.now(), + }; + }, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Task ${ctx.currentTaskIndex + 1} completed: ${event.data.result}` + ), + }), + }, + onError: { + target: 'AwaitingUserAction', + actions: assign({ + currentPlan: (ctx, event) => { + if (!ctx.currentPlan) { + return ctx.currentPlan; + } + return { + ...ctx.currentPlan, + tasks: ctx.currentPlan.tasks.map((task, index) => + index === ctx.currentTaskIndex + ? { + ...task, + status: 'failed', + error: event.data?.message || 'Task execution failed', + } + : task + ), + updatedAt: Date.now(), + }; + }, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Task ${ctx.currentTaskIndex + 1} failed: ${event.data?.message || 'Unknown error'}` + ), + }), + }, + }, + on: { + [AIChatMachineEventType.ASK_QUESTION]: { + target: 'QuestioningUser', + actions: assign({ + currentQuestion: (_ctx, event) => event.payload.question, + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'assistant', event.payload.question.question), + }), + }, + }, + }, + AwaitingUserAction: { + entry: 'saveChatState', + on: { + [AIChatMachineEventType.CONTINUE_TO_NEXT]: [ + { + cond: (ctx) => + ctx.currentPlan !== undefined && + ctx.currentTaskIndex < ctx.currentPlan.tasks.length - 1, + target: 'ExecutingTask', + actions: assign({ + currentTaskIndex: (ctx) => ctx.currentTaskIndex + 1, + chatHistory: (ctx) => + addChatMessage( + ctx.chatHistory, + 'user', + `Continue to task ${ctx.currentTaskIndex + 2}` + ), + }), + }, + { + target: 'Completed', + actions: assign({ + chatHistory: (ctx) => + addChatMessage(ctx.chatHistory, 'system', 'All tasks completed successfully!'), + }), + }, + ], + [AIChatMachineEventType.FINETUNE_TASK]: { + target: 'FinetuningTask', + actions: assign({ + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', `Finetune request: ${event.payload.instructions}`), + }), + }, + }, + }, + FinetuningTask: { + invoke: { + id: 'finetuneTask', + src: 'finetuneTask', + onDone: { + target: 'AwaitingUserAction', + actions: assign({ + currentPlan: (ctx, event) => { + if (!ctx.currentPlan) { + return ctx.currentPlan; + } + return { + ...ctx.currentPlan, + tasks: ctx.currentPlan.tasks.map((task, index) => + index === ctx.currentTaskIndex + ? { ...task, result: event.data.result, status: 'completed' } + : task + ), + updatedAt: Date.now(), + }; + }, + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Task ${ctx.currentTaskIndex + 1} refined: ${event.data.result}` + ), + }), + }, + onError: { + target: 'AwaitingUserAction', + actions: assign({ + errorMessage: (_ctx, event) => event.data?.message || 'Failed to finetune task', + chatHistory: (ctx, event) => + addChatMessage( + ctx.chatHistory, + 'assistant', + `Finetuning failed: ${event.data?.message || 'Unknown error'}` + ), + }), + }, + }, + }, + QuestioningUser: { + entry: 'saveChatState', + on: { + [AIChatMachineEventType.ANSWER_QUESTION]: { + target: 'ExecutingTask', + actions: assign({ + currentQuestion: (_ctx) => undefined, + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', event.payload.answer), + }), + }, + }, + }, + Completed: { + entry: 'saveChatState', + on: { + [AIChatMachineEventType.SUBMIT_PROMPT]: { + target: 'CreatingPlan', + actions: assign({ + initialPrompt: (_ctx, event) => event.payload.prompt, + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), + currentPlan: (_ctx) => undefined, + currentTaskIndex: (_ctx) => -1, + errorMessage: (_ctx) => undefined, + }), + }, + }, + }, + Error: { + on: { + [AIChatMachineEventType.RETRY]: [ + { + cond: (ctx) => ctx.currentPlan !== undefined, + target: 'PlanReview', + actions: assign({ + errorMessage: (_ctx) => undefined, + }), + }, + { + target: 'Idle', + actions: assign({ + errorMessage: (_ctx) => undefined, + }), + }, + ], + [AIChatMachineEventType.RESET]: { + target: 'Idle', + }, + }, + }, + }, +}); + +// Service implementations +const createPlanService = async (context: AIChatMachineContext, event: any): Promise<{ plan: Plan }> => { + const requestBody: GenerateCodeRequest = { + usecase: event.payload.prompt, + chatHistory: event.payload.chatHistory, + operationType: "CODE_GENERATION", + fileAttachmentContents: [], + }; + await generatDesign(requestBody); + + const tasks: Task[] = [ + { + id: generateId(), + description: 'Task 1: Analyze requirements', + status: 'pending', + }, + { + id: generateId(), + description: 'Task 2: Design solution', + status: 'pending', + }, + { + id: generateId(), + description: 'Task 3: Implement code', + status: 'pending', + }, + ]; + + const plan: Plan = { + id: generateId(), + tasks, + createdAt: Date.now(), + updatedAt: Date.now(), + }; + + return { plan }; +}; + +const updatePlanService = async (context: AIChatMachineContext, event: any): Promise<{ plan: Plan }> => { + // TODO: Implement actual plan update logic using AI + // This is a placeholder implementation + await new Promise((resolve) => setTimeout(resolve, 1000)); + + if (!context.currentPlan) { + throw new Error('No plan to update'); + } + + const updatedPlan: Plan = { + ...context.currentPlan, + updatedAt: Date.now(), + }; + + return { plan: updatedPlan }; +}; + +const executeTaskService = async (context: AIChatMachineContext, event: any): Promise<{ result: string }> => { + // TODO: Implement actual task execution logic + // This is a placeholder implementation + await new Promise((resolve) => setTimeout(resolve, 2000)); + + if (!context.currentPlan || context.currentTaskIndex < 0) { + throw new Error('No task to execute'); + } + + const task = context.currentPlan.tasks[context.currentTaskIndex]; + return { result: `Completed: ${task.description}` }; +}; + +const finetuneTaskService = async (context: AIChatMachineContext, event: any): Promise<{ result: string }> => { + // TODO: Implement actual task finetuning logic + // This is a placeholder implementation + await new Promise((resolve) => setTimeout(resolve, 1500)); + + if (!context.currentPlan || context.currentTaskIndex < 0) { + throw new Error('No task to finetune'); + } + + const task = context.currentPlan.tasks[context.currentTaskIndex]; + return { result: `Refined: ${task.description} with instructions: ${event.payload.instructions}` }; +}; + +// State persistence functions +const CHAT_STATE_STORAGE_KEY_PREFIX = 'ballerina.ai.chat.state'; + +/** + * Gets the storage key for the current project + * @param projectId The project identifier + * @returns The storage key for this project + */ +const getStorageKey = (projectId: string): string => { + return `${CHAT_STATE_STORAGE_KEY_PREFIX}.${projectId}`; +}; + +/** + * Saves the chat state for the current project + * @param context The chat machine context + */ +const saveChatState = (context: AIChatMachineContext) => { + try { + if (!context.projectId) { + console.warn('No project ID available, skipping state save'); + return; + } + + const stateToSave = { + initialPrompt: context.initialPrompt, + chatHistory: context.chatHistory, + currentPlan: context.currentPlan, + currentTaskIndex: context.currentTaskIndex, + sessionId: context.sessionId, + projectId: context.projectId, + savedAt: Date.now(), + }; + + const storageKey = getStorageKey(context.projectId); + extension.context?.globalState.update(storageKey, stateToSave); + + // Also save a list of all project IDs for management purposes + const allProjectIds = extension.context?.globalState.get(`${CHAT_STATE_STORAGE_KEY_PREFIX}.projects`) || []; + if (!allProjectIds.includes(context.projectId)) { + allProjectIds.push(context.projectId); + extension.context?.globalState.update(`${CHAT_STATE_STORAGE_KEY_PREFIX}.projects`, allProjectIds); + } + } catch (error) { + console.error('Failed to save chat state:', error); + } +}; + +/** + * Clears the chat state for a specific project + * @param context The chat machine context + */ +const clearChatStateAction = (context: AIChatMachineContext) => { + try { + if (!context.projectId) { + console.warn('No project ID available, skipping state clear'); + return; + } + + const storageKey = getStorageKey(context.projectId); + extension.context?.globalState.update(storageKey, undefined); + console.log(`Cleared chat state for project: ${context.projectId}`); + } catch (error) { + console.error('Failed to clear chat state:', error); + } +}; + +/** + * Loads the chat state for the current project + * @param projectId Optional project ID. If not provided, uses current workspace + * @returns The saved chat state or undefined + */ +export const loadChatState = async (projectId?: string): Promise => { + try { + const targetProjectId = projectId || generateProjectId(); + const storageKey = getStorageKey(targetProjectId); + const savedState = extension.context?.globalState.get(storageKey); + + if (savedState) { + console.log(`Loaded chat state for project: ${targetProjectId}, saved at: ${savedState.savedAt ? new Date(savedState.savedAt).toISOString() : 'unknown'}`); + } + + return savedState; + } catch (error) { + console.error('Failed to load chat state:', error); + return undefined; + } +}; + +/** + * Clears the chat state for a specific project or current project + * @param projectId Optional project ID. If not provided, uses current workspace + */ +export const clearChatState = async (projectId?: string): Promise => { + try { + const targetProjectId = projectId || generateProjectId(); + const storageKey = getStorageKey(targetProjectId); + await extension.context?.globalState.update(storageKey, undefined); + console.log(`Cleared chat state for project: ${targetProjectId}`); + } catch (error) { + console.error('Failed to clear chat state:', error); + } +}; + +/** + * Gets all project IDs that have saved chat states + * @returns Array of project IDs + */ +export const getAllProjectIds = async (): Promise => { + try { + return extension.context?.globalState.get(`${CHAT_STATE_STORAGE_KEY_PREFIX}.projects`) || []; + } catch (error) { + console.error('Failed to get project IDs:', error); + return []; + } +}; + +/** + * Clears all chat states for all projects + */ +export const clearAllChatStates = async (): Promise => { + try { + const projectIds = await getAllProjectIds(); + + for (const projectId of projectIds) { + await clearChatState(projectId); + } + + // Clear the projects list + await extension.context?.globalState.update(`${CHAT_STATE_STORAGE_KEY_PREFIX}.projects`, []); + console.log('Cleared all chat states'); + } catch (error) { + console.error('Failed to clear all chat states:', error); + } +}; + +/** + * Gets metadata about saved chat states + * @returns Array of project metadata + */ +export const getChatStateMetadata = async (): Promise> => { + try { + const projectIds = await getAllProjectIds(); + const metadata = []; + + for (const projectId of projectIds) { + const state = await loadChatState(projectId); + if (state) { + const savedState = state as AIChatMachineContext & { savedAt?: number }; + metadata.push({ + projectId, + savedAt: savedState.savedAt, + sessionId: savedState.sessionId, + taskCount: savedState.currentPlan?.tasks.length || 0, + }); + } + } + + return metadata; + } catch (error) { + console.error('Failed to get chat state metadata:', error); + return []; + } +}; + +// Create and export the state machine service +const chatStateService = interpret( + chatMachine.withConfig({ + services: { + createPlan: createPlanService, + updatePlan: updatePlanService, + executeTask: executeTaskService, + finetuneTask: finetuneTaskService, + }, + actions: { + saveChatState: (context) => saveChatState(context), + clearChatState: (context) => clearChatStateAction(context), + }, + }) +); + +const isExtendedEvent = ( + arg: K | AIChatMachineSendableEvent +): arg is Extract => { + return typeof arg !== 'string'; +}; + +export const AIChatStateMachine = { + initialize: () => { + chatStateService.start(); + + // Attempt to restore state on initialization for current project + const projectId = generateProjectId(); + loadChatState(projectId).then((savedState) => { + if (savedState && savedState.sessionId && savedState.projectId === projectId) { + console.log(`Restoring chat state for project: ${projectId}`); + chatStateService.send({ + type: AIChatMachineEventType.RESTORE_STATE, + payload: { state: savedState }, + }); + } else { + console.log(`No saved state found for project: ${projectId}, starting fresh`); + } + }).catch((error) => { + console.error('Failed to restore chat state:', error); + }); + }, + service: () => chatStateService, + context: () => chatStateService.getSnapshot().context, + state: () => chatStateService.getSnapshot().value as AIChatMachineStateValue, + sendEvent: ( + event: K | Extract + ) => { + if (isExtendedEvent(event)) { + chatStateService.send(event as AIChatMachineSendableEvent); + } else { + chatStateService.send({ type: event } as AIChatMachineSendableEvent); + } + }, + dispose: () => { + // Save state before disposing + const context = chatStateService.getSnapshot().context; + saveChatState(context); + chatStateService.stop(); + }, + /** + * Gets the current project ID + */ + getProjectId: () => { + return chatStateService.getSnapshot().context.projectId || generateProjectId(); + }, + /** + * Manually saves the current state + */ + saveState: () => { + const context = chatStateService.getSnapshot().context; + saveChatState(context); + }, + /** + * Clears the current project's chat history + */ + clearHistory: async () => { + const projectId = chatStateService.getSnapshot().context.projectId; + if (projectId) { + await clearChatState(projectId); + } + // Send reset event to clear in-memory state + chatStateService.send({ type: AIChatMachineEventType.RESET }); + }, +}; diff --git a/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts b/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts index 0ad25d5aaa1..7b344b5eaa9 100644 --- a/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts +++ b/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts @@ -51,7 +51,12 @@ import { currentThemeChanged, ChatNotify, onChatNotify, - AIMachineSendableEvent + AIMachineSendableEvent, + sendAIChatStateEvent, + AIChatMachineEventType, + aiChatStateChanged, + AIChatMachineSendableEvent, + AIChatMachineStateValue } from "@wso2/ballerina-core"; import { LangClientRpcClient } from "./rpc-clients/lang-client/rpc-client"; import { LibraryBrowserRpcClient } from "./rpc-clients/library-browser/rpc-client"; @@ -195,6 +200,14 @@ export class BallerinaRpcClient { this.messenger.sendRequest(sendAIStateEvent, HOST_EXTENSION, event); } + onAIChatStateChanged(callback: (state: AIChatMachineStateValue) => void) { + this.messenger.onNotification(aiChatStateChanged, callback); + } + + sendAIChatStateEvent(event: AIChatMachineEventType | AIChatMachineSendableEvent) { + this.messenger.sendRequest(sendAIChatStateEvent, HOST_EXTENSION, event); + } + onProjectContentUpdated(callback: (state: boolean) => void) { this.messenger.onNotification(projectContentUpdated, callback); } 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 98fc0b30f38..89dbd0f79f5 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 @@ -48,6 +48,8 @@ import { DocGenerationRequest, DocGenerationType, FileChanges, + AIChatMachineEventType, + AIChatMachineStateValue, } from "@wso2/ballerina-core"; import { useRpcContext } from "@wso2/ballerina-rpc-client"; @@ -149,6 +151,7 @@ const AIChat: React.FC = () => { const [isAddingToWorkspace, setIsAddingToWorkspace] = useState(false); const [showSettings, setShowSettings] = useState(false); + const [aiChatStateMachineState, setAiChatStateMachineState] = useState("Idle"); //TODO: Need a better way of storing data related to last generation to be in the repair state. const currentDiagnosticsRef = useRef([]); @@ -205,8 +208,8 @@ const AIChat: React.FC = () => { chatLocation = (await rpcClient.getVisualizerLocation()).projectUri; setIsReqFileExists( chatLocation != null && - chatLocation != undefined && - (await rpcClient.getAiPanelRpcClient().isRequirementsSpecificationFileExist(chatLocation)) + chatLocation != undefined && + (await rpcClient.getAiPanelRpcClient().isRequirementsSpecificationFileExist(chatLocation)) ); generateNaturalProgrammingTemplate(isReqFileExists); @@ -274,6 +277,10 @@ const AIChat: React.FC = () => { }); }, []); + rpcClient?.onAIChatStateChanged((newState: AIChatMachineStateValue) => { + setAiChatStateMachineState(newState); + }); + rpcClient?.onChatNotify((response: ChatNotify) => { // TODO: Need to handle the content as step blocks const type = response.type; @@ -1678,14 +1685,10 @@ const AIChat: React.FC = () => { } async function processDesignGeneration(useCase: string, message: string) { - const requestBody: GenerateCodeRequest = { - usecase: useCase, - chatHistory: chatArray, - operationType: CodeGenerationType.CODE_GENERATION, - fileAttachmentContents: [], - }; - - await rpcClient.getAiPanelRpcClient().generateDesign(requestBody); + rpcClient.sendAIChatStateEvent({ + type: AIChatMachineEventType.SUBMIT_PROMPT, + payload: { prompt: useCase } + }); } async function handleStop() { @@ -1715,6 +1718,8 @@ const AIChat: React.FC = () => { setMessages((prevMessages) => []); localStorage.removeItem(`chatArray-AIGenerationChat-${projectUuid}`); + + rpcClient.sendAIChatStateEvent(AIChatMachineEventType.RESET); } const questionMessages = messages.filter((message) => message.type === "question"); @@ -1927,6 +1932,7 @@ const AIChat: React.FC = () => {
{/* {`Resets in: 30 days`} */} +
State: {aiChatStateMachineState}
+ {activeTasks && activeTasks.length > 0 && ( + + + + )}
{Array.isArray(otherMessages) && otherMessages.length === 0 && ( @@ -2045,6 +2211,10 @@ const AIChat: React.FC = () => { failed={segment.failed} /> ); + } else if (segment.type === SegmentType.Todo) { + // Skip rendering todo in chat messages - tasks are shown only in top panel + // The tags remain in message content for agent context + return null; } else if (segment.type === SegmentType.Attachment) { return ( @@ -2205,6 +2375,16 @@ const AIChat: React.FC = () => { )} {showSettings && setShowSettings(false)}>} + {approvalRequest && ( + + )} ); }; @@ -2328,6 +2508,7 @@ export enum SegmentType { Text = "Text", Progress = "Progress", ToolCall = "ToolCall", + Todo = "Todo", Attachment = "Attachment", InlineCode = "InlineCode", References = "References", @@ -2403,7 +2584,7 @@ export function splitContent(content: string): Segment[] { // Combined regex to capture either ``` code ``` or Text const regex = - /\s*```(\w+)\s*([\s\S]*?)```\s*<\/code>|([\s\S]*?)<\/progress>|([\s\S]*?)<\/toolcall>|([\s\S]*?)<\/attachment>|([\s\S]*?)<\/scenario>|([\s\S]*?)<\/button>|([\s\S]*?)|([\s\S]*?)/g; + /\s*```(\w+)\s*([\s\S]*?)```\s*<\/code>|([\s\S]*?)<\/progress>|([\s\S]*?)<\/toolcall>|([\s\S]*?)<\/todo>|([\s\S]*?)<\/attachment>|([\s\S]*?)<\/scenario>|([\s\S]*?)<\/button>|([\s\S]*?)|([\s\S]*?)/g; let match; let lastIndex = 0; @@ -2460,8 +2641,26 @@ export function splitContent(content: string): Segment[] { text: toolcallText, }); } else if (match[7]) { + // block matched + const todoData = match[7]; + + updateLastProgressSegmentLoading(); + try { + const parsedData = JSON.parse(todoData); + segments.push({ + type: SegmentType.Todo, + loading: false, + text: "", + tasks: parsedData.tasks || [], + message: parsedData.message || "" + }); + } catch (error) { + // If parsing fails, show as text + console.error("Failed to parse todo data:", error); + } + } else if (match[8]) { // block matched - const attachmentName = match[7].trim(); + const attachmentName = match[8].trim(); updateLastProgressSegmentLoading(); @@ -2476,9 +2675,9 @@ export function splitContent(content: string): Segment[] { text: attachmentName, }); } - } else if (match[8]) { + } else if (match[9]) { // block matched - const scenarioContent = match[8].trim(); + const scenarioContent = match[9].trim(); updateLastProgressSegmentLoading(true); segments.push({ @@ -2486,10 +2685,10 @@ export function splitContent(content: string): Segment[] { loading: false, text: scenarioContent, }); - } else if (match[9]) { + } else if (match[10]) { // + + + ) : ( + <> + + + + )} + + + + ); +}; + +export default ApprovalDialog; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx new file mode 100644 index 00000000000..a0a6dcc7956 --- /dev/null +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com) All Rights Reserved. + * + * WSO2 LLC. licenses this file to you under the Apache License, + * Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { keyframes } from "@emotion/css"; +import styled from "@emotion/styled"; +import React, { useState } from "react"; + +const spin = keyframes` + from { transform: rotate(0deg); } + to { transform: rotate(360deg); } +`; + +const TodoContainer = styled.div` + background-color: transparent; + border: none; + border-radius: 0; + padding: 0; + margin: 0; + font-family: var(--vscode-editor-font-family); + font-size: 13px; + color: var(--vscode-editor-foreground); + min-height: 32px; +`; + +const TodoHeader = styled.div<{ clickable?: boolean }>` + font-weight: 600; + margin-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '12px'}; + padding-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '8px'}; + border-bottom: ${(props: { clickable?: boolean }) => + props.clickable ? 'none' : '1px solid var(--vscode-panel-border)'}; + display: flex; + align-items: center; + gap: 8px; + cursor: ${(props: { clickable?: boolean }) => props.clickable ? 'pointer' : 'default'}; + user-select: none; + padding: 4px; + border-radius: 4px; + + &:hover { + background-color: ${(props: { clickable?: boolean }) => + props.clickable ? 'var(--vscode-list-hoverBackground)' : 'transparent'}; + } +`; + +const ChevronIcon = styled.span<{ expanded: boolean }>` + transition: transform 0.2s ease; + transform: ${(props: { expanded: boolean }) => props.expanded ? 'rotate(90deg)' : 'rotate(0deg)'}; + display: flex; + align-items: center; +`; + +const MinimalTaskInfo = styled.span` + color: var(--vscode-descriptionForeground); + font-weight: 400; + font-size: 12px; + margin-left: 4px; +`; + +const TodoList = styled.div` + display: flex; + flex-direction: column; + gap: 8px; +`; + +const TodoItem = styled.div<{ status: string }>` + display: flex; + align-items: flex-start; + gap: 10px; + padding: 8px; + border-radius: 4px; + background-color: ${(props: { status: string }) => + props.status === "completed" + ? "var(--vscode-list-hoverBackground)" + : "transparent"}; + opacity: ${(props: { status: string }) => (props.status === "completed" ? 0.7 : 1)}; + transition: all 0.2s ease; +`; + +const TodoIcon = styled.span<{ status: string }>` + flex-shrink: 0; + display: flex; + align-items: center; + justify-content: center; + width: 18px; + height: 18px; + margin-top: 2px; + + &.pending { + .codicon { + color: var(--vscode-descriptionForeground); + } + } + + &.in_progress { + .codicon { + color: var(--vscode-charts-blue); + animation: ${spin} 1s linear infinite; + } + } + + &.completed { + .codicon { + color: var(--vscode-testing-iconPassed); + } + } +`; + +const TodoText = styled.div<{ status: string }>` + flex: 1; + text-decoration: ${(props: { status: string }) => + props.status === "completed" ? "line-through" : "none"}; +`; + +const TodoNumber = styled.span` + color: var(--vscode-descriptionForeground); + font-weight: 600; + margin-right: 4px; +`; + +export interface Task { + id: string; + description: string; + status: "pending" | "in_progress" | "completed"; +} + +interface TodoSectionProps { + tasks: Task[]; + message?: string; +} + +const getStatusIcon = (status: string): { className: string; icon: string } => { + switch (status) { + case "in_progress": + return { className: "in_progress", icon: "codicon-sync" }; + case "completed": + return { className: "completed", icon: "codicon-check" }; + case "pending": + default: + return { className: "pending", icon: "codicon-circle-outline" }; + } +}; + +const TodoSection: React.FC = ({ tasks, message }) => { + const [isExpanded, setIsExpanded] = useState(true); + const completedCount = tasks.filter((t) => t.status === "completed").length; + const inProgressTask = tasks.find((t) => t.status === "in_progress"); + const allCompleted = completedCount === tasks.length; + const hasInProgress = !!inProgressTask; + + const toggleExpanded = () => { + setIsExpanded(!isExpanded); + }; + + // Determine status text + const getStatusText = () => { + if (allCompleted) return "completed"; + if (hasInProgress) return "in progress"; + return "ongoing"; + }; + + return ( + + + + + + + + Implementation Tasks ({completedCount}/{tasks.length}{" "} + {getStatusText()}) + + {!isExpanded && inProgressTask && ( + + > {inProgressTask.description} + + )} + + {isExpanded && ( + <> + {message && ( +
+ {message} +
+ )} + + {tasks.map((task, index) => { + const statusInfo = getStatusIcon(task.status); + return ( + + + + + + {index + 1}. + {task.description} + + + ); + })} + + + )} + + ); +}; + +export default TodoSection; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx index f8ef0d9f97d..808f4c3c0cf 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx @@ -45,6 +45,15 @@ export const HeaderButtons = styled.div({ marginRight: "10px", }); +export const TodoPanel = styled.div` + max-height: 30vh; + overflow-y: auto; + border-bottom: 2px solid var(--vscode-panel-border); + background-color: var(--vscode-editor-background); + padding: 12px 20px; + flex-shrink: 0; +`; + export const Main = styled.main({ flex: 1, flexDirection: "column", From 49ced2c33c638bd7bfd12bae05e6fef81239a344 Mon Sep 17 00:00:00 2001 From: Anjana Supun Date: Tue, 21 Oct 2025 14:34:20 +0530 Subject: [PATCH 004/619] Initial integration of /design with state machine --- common/config/rush/pnpm-lock.yaml | 3 + .../ballerina-core/src/state-machine-types.ts | 25 ++++++- .../src/features/ai/service/design/design.ts | 73 ++++++------------- .../ai/service/libs/task_write_tool.ts | 71 +++++++----------- .../src/rpc-managers/ai-panel/rpc-manager.ts | 4 +- .../src/views/ai-panel/aiChatMachine.ts | 32 ++------ .../views/AIPanel/components/AIChat/index.tsx | 3 +- 7 files changed, 85 insertions(+), 126 deletions(-) diff --git a/common/config/rush/pnpm-lock.yaml b/common/config/rush/pnpm-lock.yaml index e483ed47f52..2fad9e583bb 100644 --- a/common/config/rush/pnpm-lock.yaml +++ b/common/config/rush/pnpm-lock.yaml @@ -24169,7 +24169,10 @@ snapshots: jest-runner: 25.5.4 jest-runtime: 25.5.4 transitivePeerDependencies: + - bufferutil + - canvas - supports-color + - utf-8-validate '@jest/test-sequencer@29.7.0': dependencies: diff --git a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts index 7618d170e87..d2961a09e41 100644 --- a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts +++ b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts @@ -371,12 +371,29 @@ export interface ChatMessage { timestamp: number; } +/** + * Task status enum + */ +export enum TaskStatus { + PENDING = "pending", + IN_PROGRESS = "in_progress", + COMPLETED = "completed" +} + +export enum TaskTypes { + SERVICE_DESIGN = "service_design", + CONNECTIONS_INIT = "connections_init", + IMPLEMENTATION = "implementation" +} + +/** + * Task interface representing a single implementation task + */ export interface Task { id: string; description: string; - status: 'pending' | 'in-progress' | 'completed' | 'failed'; - result?: string; - error?: string; + status: TaskStatus; + type : TaskTypes; } export interface Plan { @@ -480,4 +497,4 @@ export const sendAIStateEvent: RequestType = { method: 'currentThemeChanged' }; export const aiChatStateChanged: NotificationType = { method: 'aiChatStateChanged' }; -export const sendAIChatStateEvent: RequestType = { method: 'sendAIChatStateEvent' }; \ No newline at end of file +export const sendAIChatStateEvent: RequestType = { method: 'sendAIChatStateEvent' }; 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 debf2af2392..35faf54fba0 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 @@ -14,18 +14,18 @@ // specific language governing permissions and limitations // under the License. -import { Command, GenerateCodeRequest } from "@wso2/ballerina-core"; +import { Command, GenerateCodeRequest, Task } from "@wso2/ballerina-core"; import { ModelMessage, stepCountIs, streamText } from "ai"; -import { getAnthropicClient, ANTHROPIC_HAIKU, getProviderCacheControl } from "../connection"; +import { getAnthropicClient, ANTHROPIC_HAIKU, getProviderCacheControl, ANTHROPIC_SONNET_4 } from "../connection"; import { getErrorMessage, populateHistory } from "../utils"; import { CopilotEventHandler, createWebviewEventHandler } from "../event"; import { AIPanelAbortController } from "../../../../rpc-managers/ai-panel/utils"; -import { createTaskWriteTool, TASK_WRITE_TOOL_NAME, resolveTaskApproval } from "../libs/task_write_tool"; +import { createTaskWriteTool, TASK_WRITE_TOOL_NAME, resolveTaskApproval, TaskWriteResult } from "../libs/task_write_tool"; // Re-export for RPC manager to use export { resolveTaskApproval as resolveApproval }; -export async function generateDesignCore(params: GenerateCodeRequest, eventHandler: CopilotEventHandler): Promise { +export async function generateDesignCore(params: GenerateCodeRequest, eventHandler: CopilotEventHandler): Promise { // Populate chat history and add user message const historyMessages = populateHistory(params.chatHistory); const cacheOptions = await getProviderCacheControl(); @@ -33,54 +33,24 @@ export async function generateDesignCore(params: GenerateCodeRequest, eventHandl const allMessages: ModelMessage[] = [ { role: "system", - content: `You are an expert assistant specializing in Ballerina code generation. You should ONLY answer Ballerina related queries. + content: `You are an expert assistant to help with writing ballerina integrations. You will be helping with designing a solution for user query in a step-by-step manner. -Your primary responsibility is to generate a high-level design for a Ballerina integration based on the user's requirements, and then implement it. +You have access to ${TASK_WRITE_TOOL_NAME} tool to create and manage tasks. This plan will be visible to the user and the execution will be guided on the tasks you create. -IMPORTANT: Before executing any other actions to fulfill the user's query, you MUST first design a comprehensive plan. - -When creating a design, provide a structured outline that includes: - -1. **Overview**: A brief summary of the integration's purpose and goals -2. **Components**: List all major components/modules needed (e.g., HTTP services, clients, data models, connectors) -3. **Data Flow**: Describe how data moves through the system - - Input sources and formats - - Transformation steps - - Output destinations and formats -4. **Interactions**: Detail the interactions between components - - API endpoints and their purposes - - External service integrations - - Database or storage interactions -5. **Error Handling**: Outline error handling strategy -6. **Security Considerations**: Note any authentication, authorization, or data security requirements - -After creating the high-level design, if the implementation requires MORE THAN THREE distinct steps, break it down into specific implementation tasks and execute them step by step. Do NOT mention internal tool names to the user - just naturally describe what you're doing (e.g., "I'll now break this down into implementation tasks" instead of "I'll use the TaskWrite tool"). - -Format your design plan clearly using markdown with appropriate headings and bullet points for readability.`, - providerOptions: cacheOptions, - }, - { - role: "system", - content: ` if you are generating code, ensure to: - - Decide which libraries need to be imported (Avoid importing lang.string, lang.boolean, lang.float, lang.decimal, lang.int, lang.map langlibs as they are already imported by default). - - Determine the necessary client initialization. - - Define Types needed for the query in the types.bal file. - - Outline the service OR main function for the query. - - Outline the required function usages as noted in Step 2. - - Based on the types of identified functions, plan the data flow. Transform data as necessary. - - Finally, provide a - Example Codeblock segment: - - \`\`\`ballerina - //code goes here - \`\`\` - +Do NOT mention internal tool names to the user - just naturally describe what you're doing (e.g., "I'll now break this down into implementation tasks" instead of "I'll use the ${TASK_WRITE_TOOL_NAME} tool"). `, + providerOptions: cacheOptions, }, ...historyMessages, { role: "user", - content: params.usecase, + content: `The first step is to create a very high level consise design plan for the following requirement. +Then you must create the tasks using ${TASK_WRITE_TOOL_NAME} tool accoringly. Then the user will approve the tasks or ask for modifications. + + + +${params.usecase} +`, }, ]; @@ -90,7 +60,7 @@ Format your design plan clearly using markdown with appropriate headings and bul }; const { fullStream } = streamText({ - model: await getAnthropicClient(ANTHROPIC_HAIKU), + model: await getAnthropicClient(ANTHROPIC_SONNET_4), maxOutputTokens: 8192, temperature: 0, messages: allMessages, @@ -99,6 +69,9 @@ Format your design plan clearly using markdown with appropriate headings and bul abortSignal: AIPanelAbortController.getInstance().signal, }); + // TODO: Will it call this tool multiple times? + let finalTasks: Task[] = []; + eventHandler({ type: "start" }); for await (const part of fullStream) { @@ -123,7 +96,7 @@ Format your design plan clearly using markdown with appropriate headings and bul // Emit tool result event with full task list if (toolName === TASK_WRITE_TOOL_NAME && result) { - const taskResult = result as any; + const taskResult = result as TaskWriteResult; eventHandler({ type: "tool_result", toolName, @@ -133,6 +106,7 @@ Format your design plan clearly using markdown with appropriate headings and bul allTasks: taskResult.tasks // Tool returns complete task list } }); + finalTasks = taskResult.tasks; } break; } @@ -151,13 +125,14 @@ Format your design plan clearly using markdown with appropriate headings and bul } } } + return finalTasks; } // Main public function that uses the default event handler -export async function generatDesign(params: GenerateCodeRequest): Promise { +export async function generateDesign(params: GenerateCodeRequest): Promise { const eventHandler = createWebviewEventHandler(Command.Design); try { - await generateDesignCore(params, eventHandler); + return await generateDesignCore(params, eventHandler); } catch (error) { console.error("Error during design generation:", error); eventHandler({ type: "error", content: getErrorMessage(error) }); 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 dbe2f391599..dc3e45fa07e 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,27 +17,10 @@ import { tool } from 'ai'; import { z } from 'zod'; import { CopilotEventHandler } from '../event'; +import { Task, TaskStatus, TaskTypes } from '@wso2/ballerina-core'; export const TASK_WRITE_TOOL_NAME = "TaskWrite"; -/** - * Task status enum - */ -export enum TaskStatus { - PENDING = "pending", - IN_PROGRESS = "in_progress", - COMPLETED = "completed" -} - -/** - * Task interface representing a single implementation task - */ -export interface Task { - id: string; - description: string; - status: TaskStatus; -} - /** * Result returned by TaskWrite tool */ @@ -50,12 +33,18 @@ export interface TaskWriteResult { /** * Zod schema for a single task input */ -const TaskInputSchema = z.object({ +export const TaskInputSchema = z.object({ id: z.string().optional().describe("Task ID (required when updating existing tasks, omit when creating new tasks)"), description: z.string().min(1).describe("Clear, actionable description of the task to be implemented"), - status: z.enum(["pending", "in_progress", "completed"]).describe("Current status of the task. Use 'pending' for new tasks, 'in_progress' when starting work, 'completed' when done.") + status: z.enum(["pending", "in_progress", "completed"]).describe("Current status of the task. Use 'pending' for new tasks, 'in_progress' when starting work, 'completed' when done."), + type: z.enum(["service_design", "connections_init", "implementation"]).describe("Type of the implementation task. service_design will only generate the http service contract. not the implementation. connections_init will only generate the connection initializations. All of the other tasks will be of type implementation.") }); +/** + * Type for a single task input + */ +export type TaskInput = z.infer; + /** * Zod schema for TaskWrite tool input */ @@ -93,8 +82,9 @@ export function resolveTaskApproval(response: { approved: boolean; comment?: str export function createTaskWriteTool(eventHandler?: CopilotEventHandler) { return tool({ description: `Create and update implementation tasks for the design plan. - -ONLY use this tool if you have MORE THAN 3 tasks. For 3 or fewer tasks, proceed directly without this tool. +## Task Ordering: +- Tasks should be ordered sequentially as they need to be executed. +- Prioritize service design, then connection initializations, then implementation tasks. ## CRITICAL RULE - ALWAYS SEND ALL TASKS: This tool is STATELESS. Every call MUST include ALL tasks. @@ -112,9 +102,9 @@ This tool is STATELESS. Every call MUST include ALL tasks. Send ALL tasks with status "pending", no IDs. Example: [ - {"description": "Create data types", "status": "pending"}, - {"description": "Implement REST API", "status": "pending"}, - {"description": "Add error handling", "status": "pending"} + {"description": "Create the HTTP service contract", "status": "pending", "type": "service_design"}, + {"description": "Create the MYSQL Connection", "status": "pending", "type": "connections_init"}, + {"description": "Implement the resource functions", "status": "pending", "type": "implementation"} ] ## UPDATING TASKS (Every Other Call): @@ -131,35 +121,25 @@ Workflow per task: Example (3 tasks total): Start task 1 - Send ALL: [ - {"id": "1", "description": "Create data types", "status": "in_progress"}, - {"id": "2", "description": "Implement REST API", "status": "pending"}, - {"id": "3", "description": "Add error handling", "status": "pending"} + {"id": "1", "description": "Create the HTTP service contract", "status": "in_progress", "type": "service_design"}, + {"id": "2", "description": "Create the MYSQL Connection", "status": "pending", "type": "connections_init"}, + {"id": "3", "description": "Implement the resource functions", "status": "pending", "type": "implementation"} ] Complete task 1 - Send ALL: [ - {"id": "1", "description": "Create data types", "status": "completed"}, - {"id": "2", "description": "Implement REST API", "status": "pending"}, - {"id": "3", "description": "Add error handling", "status": "pending"} + {"id": "1", "description": "Create the HTTP service contract", "status": "complete", "type": "service_design"}, + {"id": "2", "description": "Create the MYSQL Connection", "status": "pending", "type": "connections_init"}, + {"id": "3", "description": "Implement the resource functions", "status": "pending", "type": "implementation"} ] Start task 2 - Send ALL: [ - {"id": "1", "description": "Create data types", "status": "completed"}, - {"id": "2", "description": "Implement REST API", "status": "in_progress"}, - {"id": "3", "description": "Add error handling", "status": "pending"} + {"id": "1", "description": "Create the HTTP service contract", "status": "complete", "type": "service_design"}, + {"id": "2", "description": "Create the MYSQL Connection", "status": "in_progress", "type": "connections_init"}, + {"id": "3", "description": "Implement the resource functions", "status": "pending", "type": "implementation"} ] -WRONG: -[{"id": "1", "status": "completed"}] // Only 1 task - WILL BE REJECTED! - -CORRECT: -[ - {"id": "1", "status": "completed"}, - {"id": "2", "status": "pending"}, - {"id": "3", "status": "pending"} -] // Sending ALL tasks - Rules: - Send ALL tasks every single call (tool will reject partial lists) - Only ONE task "in_progress" at a time @@ -174,7 +154,8 @@ Rules: const allTasks: Task[] = input.tasks.map(task => ({ id: task.id || `task_${Date.now()}_${Math.random()}`, // Generate ID only if not provided description: task.description, - status: task.status as TaskStatus + status: task.status as TaskStatus, + type: task.type as TaskTypes })); console.log(`[TaskWrite Tool] Received ${allTasks.length} task(s)`); diff --git a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts index 1e43904c39d..1644fb1e412 100644 --- a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts +++ b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts @@ -120,7 +120,7 @@ import { import { attemptRepairProject, checkProjectDiagnostics } from "./repair-utils"; import { AIPanelAbortController, addToIntegration, cleanDiagnosticMessages, handleStop, isErrorCode, processMappings, requirementsSpecification, searchDocumentation } from "./utils"; import { fetchData } from "./utils/fetch-data-utils"; -import { generatDesign, generateDesignCore, resolveApproval } from "../../../src/features/ai/service/design/design"; +import { generateDesign, generateDesignCore, resolveApproval } from "../../../src/features/ai/service/design/design"; export class AiPanelRpcManager implements AIPanelAPI { @@ -1044,7 +1044,7 @@ export class AiPanelRpcManager implements AIPanelAPI { } async generateDesign(params: GenerateCodeRequest): Promise { - await generatDesign(params); + await generateDesign(params); return true; } diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts index 9c93f4f07b8..cece780dbd1 100644 --- a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts +++ b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts @@ -20,10 +20,10 @@ import { createMachine, assign, interpret } from 'xstate'; import { extension } from '../../BalExtensionContext'; import * as crypto from 'crypto'; -import { AIChatMachineContext, AIChatMachineEventType, AIChatMachineSendableEvent, AIChatMachineStateValue, ChatMessage, Plan, Task } from '@wso2/ballerina-core/lib/state-machine-types'; +import { AIChatMachineContext, AIChatMachineEventType, AIChatMachineSendableEvent, AIChatMachineStateValue, ChatMessage, Plan, Task, TaskStatus } from '@wso2/ballerina-core/lib/state-machine-types'; import { workspace } from 'vscode'; import { GenerateCodeRequest } from '@wso2/ballerina-core/lib/rpc-types/ai-panel/interfaces'; -import { generatDesign } from '../../features/ai/service/design/design'; +import { generateDesign } from '../../features/ai/service/design/design'; const generateId = () => `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; @@ -69,6 +69,7 @@ const addChatMessage = ( }; const chatMachine = createMachine({ + /** @xstate-layout N4IgpgJg5mDOIC5QCMCGAbdYBOBLAdqgLSq5EDGAFqgC4B0AwtmLQVAArqr4DEEA9vjB0CAN34BrYeWa0wnbgG0ADAF1EoAA79YuGrkEaQAD0QBmACxm6ZgJzKAbACZbTpwEYnZ97YA0IAE9EAHZg9zoADk8HZQBWW2DHC1iLAF9U-zRMHAJiUgpqeiYWfXwOLl4cbH5sOk0uGgAzGoBbOhkS+QqVdSQQbV19Qz7TBEtrO0cXN09vP0DECIdYuidl4KcLdwdgiPtY9MyMLDxCEjIqWjoAVU0IVjKFXgEhEXxxKToAVzu5J56jAM9AZ8EZRuMbPZnK4PF4fP4gghYhELHQoWYHLZYsEHO5YrEzIcQFkTrlzgUrrd7qVytweFUanUGs1sG0ftSuko1ICdMDhqBwVZIVMYbN4QsEDsIjZ4hYIqFgpZlMo0hlicccmd8pd6ABRYxgchfGkAFVQsAkfEEwjEkmEYANRpoYDNFoBfSBQ1BI3MQsm0JmcPmiKidAstgjtnD9lcZkcRJJmryF0KdH1huNbFdloZtXqtBZbQdGed2fdWl5XrBvomUOmsLmCMQ7mCqKx3nc3giUQcUWCCY1p2TFPoADECGAaF98FnzZaXjb3na6I0J1OhGXuR7KyDq2M-XXRUGm2NcZF4pYsRE1sjVUdskPyTq6OOhOvZxb6dhqnnma0V2u04unO5b9Du-ImDWwoBg24qIms0oWFM3hmBEKSdoSRL4PwEBwEYiaPtqhQ8oMu4+ggRAOCeRArJGdH0RGmH3qSWoplcACSEBYCRfLegKzadqs7jyrE7jbE42JiSe4ZOHQnZONetjXpiDjJAcaoEWSRFXMUDy0nxYGkRBoy9qs2IRHEBIOM4wS2GYJ67KiBK2C2qmxMomwEgOD5aWx9BPAASmAoi4GAADuPFVuR6E2FYOIRGYtmOHBiDxIh1lOIqqH2N2d7qj5rEjjcvw0k8kVkfxSK2GiUwWBYmVbBsYQnvEwR0A4ZguEhlnBPi7jeSxw7PumTofhI5XGYgyR0M5cTOMJUaOPZEpoTNaXXiiHU+FYA1Jk+qYAIJhaQNLXLAOAHeQUUVkZBmjDF4zxYlUIpQgwmyRYtm9rKdgdfiu2EX5L6ATOZTZhNd1TU4J4KeE8m9YlWzuHE-YaYOvlFQAil8cBemwZ04BDe4ojYnUeTsLm2B1jgnqhDjtZ1Sy-UpVNmEx+WDftOn8C09STpARPkSTbNOOTtk+NTVErfYdDKHY7j1XGiqdupzF7dperfjUguVQ9cW9s9yXBiEyhtciMRs-VN7xukqRAA */ id: 'ballerina-ai-chat', initial: 'Idle', predictableActionArguments: true, @@ -244,7 +245,7 @@ const chatMachine = createMachine index === ctx.currentTaskIndex - ? { ...task, status: 'completed', result: event.data.result } + ? { ...task, status: TaskStatus.COMPLETED, result: event.data.result } : task ), updatedAt: Date.now(), @@ -258,6 +259,7 @@ const chatMachine = createMachine index === ctx.currentTaskIndex - ? { ...task, result: event.data.result, status: 'completed' } + ? { ...task, result: event.data.result, status: TaskStatus.COMPLETED } : task ), updatedAt: Date.now(), @@ -440,25 +442,7 @@ const createPlanService = async (context: AIChatMachineContext, event: any): Pro operationType: "CODE_GENERATION", fileAttachmentContents: [], }; - await generatDesign(requestBody); - - const tasks: Task[] = [ - { - id: generateId(), - description: 'Task 1: Analyze requirements', - status: 'pending', - }, - { - id: generateId(), - description: 'Task 2: Design solution', - status: 'pending', - }, - { - id: generateId(), - description: 'Task 3: Implement code', - status: 'pending', - }, - ]; + const tasks: Task[] = await generateDesign(requestBody); const plan: Plan = { id: generateId(), 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 22f18c87eee..86e72dc7be3 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 @@ -683,8 +683,7 @@ const AIChat: React.FC = () => { if (parsedInput && "type" in parsedInput && parsedInput.type === "error") { throw new Error(parsedInput.message); } else if ("text" in parsedInput && !("command" in parsedInput)) { - await processDesignGeneration(parsedInput.text, inputText); - //await processCodeGeneration([parsedInput.text, attachments, CodeGenerationType.CODE_GENERATION], inputText); + await processCodeGeneration([parsedInput.text, attachments, CodeGenerationType.CODE_GENERATION], inputText); } else if ("command" in parsedInput) { switch (parsedInput.command) { case Command.NaturalProgramming: { From e40e08b0089ea68baa470f62626ddb0ddcf03745 Mon Sep 17 00:00:00 2001 From: RNViththagan Date: Tue, 21 Oct 2025 14:48:38 +0530 Subject: [PATCH 005/619] Improves AI design generation and task management Refactors the AI design generation process by introducing a skeleton-first approach with task management. This change enhances the AI's ability to handle complex integration designs by breaking them down into smaller, manageable tasks. It also improves the user experience by providing better visibility into the design process and allows for more granular control over the AI's actions. --- .../src/features/ai/service/design/design.ts | 104 ++++++++++++++++-- .../ai/service/libs/task_write_tool.ts | 47 ++------ .../data/commandTemplates.const.ts | 12 +- .../views/AIPanel/components/AIChat/index.tsx | 5 +- .../AIPanel/components/ApprovalDialog.tsx | 53 ++++++++- .../views/AIPanel/components/TodoSection.tsx | 102 +++++++++++++---- 6 files changed, 243 insertions(+), 80 deletions(-) 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 35faf54fba0..eb38d523ea8 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 @@ -35,19 +35,105 @@ export async function generateDesignCore(params: GenerateCodeRequest, eventHandl role: "system", content: `You are an expert assistant to help with writing ballerina integrations. You will be helping with designing a solution for user query in a step-by-step manner. +ONLY answer Ballerina-related queries. + +# Plan Mode Approach + +Follow the skeleton-first approach for all tasks: + +## Step 1: Create High-Level Design +Create a comprehensive design plan with: + +**1. Overview** +- Brief summary of the integration's purpose and goals + +**2. Components & Architecture** +- Data types and models needed +- HTTP services, clients, or main function structure +- External connectors and integrations + +**3. Implementation Approach - Skeleton-First Strategy** +ALWAYS follow this order for complex implementations: +- First: Define skeleton (types, function signatures, service structure) +- Second: Set up connections (clients, endpoints, configurations) +- Third: Implement business logic (data flow, transformations, error handling) +- Fourth: Add security and final touches + +**4. Data Flow** +- Input sources and formats +- Transformation steps +- Output destinations and formats + +## Step 2: Break Down Into Tasks and Execute + +You MUST use task management to implement the skeleton-first approach systematically. + +**REQUIRED: Use Task Management** You have access to ${TASK_WRITE_TOOL_NAME} tool to create and manage tasks. This plan will be visible to the user and the execution will be guided on the tasks you create. -Do NOT mention internal tool names to the user - just naturally describe what you're doing (e.g., "I'll now break this down into implementation tasks" instead of "I'll use the ${TASK_WRITE_TOOL_NAME} tool"). +- Break down the implementation into specific, actionable tasks following the skeleton-first order +- Track each task as you work through them +- Mark tasks as you start and complete them +- This ensures you don't miss critical steps + +**Task Breakdown Example (Skeleton-First)**: +1. Define data types and record structures (Skeleton) +2. Create service/function signatures (Skeleton) +3. Initialize HTTP clients and connections (Connections) +4. Implement main business logic (Implementation) + +**Critical**: +- Task management is MANDATORY for all implementations +- It prevents missing steps and ensures systematic implementation +- Users get visibility into your progress +- Do NOT mention internal tool names to the user - just naturally describe what you're doing (e.g., "I'll now break this down into implementation tasks" instead of "I'll use the ${TASK_WRITE_TOOL_NAME} tool") + +**Execution Flow**: +1. Create high-level design plan and break it into tasks using ${TASK_WRITE_TOOL_NAME} +2. The tool will wait for PLAN APPROVAL from the user +3. Once plan is APPROVED (success: true in tool response), IMMEDIATELY start the execution cycle: + + **For each task:** + - Mark task as in_progress using ${TASK_WRITE_TOOL_NAME} (send ALL tasks) + - Implement the task completely (write the Ballerina code) + - Mark task as completed using ${TASK_WRITE_TOOL_NAME} (send ALL tasks) + - The tool will wait for TASK COMPLETION APPROVAL from the user + - Once approved (success: true), immediately start the next task + - Repeat until ALL tasks are done + +4. **Critical**: After each approval (both plan and task completions), immediately proceed to the next step without any delay or additional prompting + +## Code Generation Guidelines + +When generating Ballerina code: + +1. **Imports**: Import required libraries + - Do NOT import these (already available by default): lang.string, lang.boolean, lang.float, lang.decimal, lang.int, lang.map + +2. **Structure**: + - Define types in types.bal file + - Initialize necessary clients + - Create service OR main function + - Plan data flow and transformations + +3. **Code Format**: + \`\`\` + + \`\`\`ballerina + // Your Ballerina code here + \`\`\` + + \`\`\` `, providerOptions: cacheOptions, }, ...historyMessages, { role: "user", - content: `The first step is to create a very high level consise design plan for the following requirement. -Then you must create the tasks using ${TASK_WRITE_TOOL_NAME} tool accoringly. Then the user will approve the tasks or ask for modifications. + content: `Create a high-level design plan for the following requirement and break it down into implementation tasks. + +After the plan is approved, execute all tasks continuously following the execution flow defined in the system prompt. - ${params.usecase} `, @@ -69,7 +155,7 @@ ${params.usecase} abortSignal: AIPanelAbortController.getInstance().signal, }); - // TODO: Will it call this tool multiple times? + // TODO: Will it call this tool multiple times? let finalTasks: Task[] = []; eventHandler({ type: "start" }); @@ -103,13 +189,17 @@ ${params.usecase} toolOutput: { success: taskResult.success, message: taskResult.message, - allTasks: taskResult.tasks // Tool returns complete task list - } + allTasks: taskResult.tasks, // Tool returns complete task list + }, }); finalTasks = taskResult.tasks; } break; } + case "text-start": { + eventHandler({ type: "content_block", content: " \n" }); + break; + } case "error": { const error = part.error; console.error("Error during design generation:", error); 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 dc3e45fa07e..50c1d454c10 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 @@ -96,7 +96,7 @@ This tool is STATELESS. Every call MUST include ALL tasks. ## USER APPROVAL REQUIRED: 1. **Plan Approval**: User approves/rejects initial task list -2. **Task Approval**: User approves/rejects each completed task +2. **Task Completion Approval**: User approves/rejects each completed task before moving to the next ## CREATING TASKS (First Call): Send ALL tasks with status "pending", no IDs. @@ -110,13 +110,13 @@ Example: ## UPDATING TASKS (Every Other Call): Send ALL tasks with IDs and updated statuses. -Workflow per task: +Workflow per task (after plan approval): 1. Mark in_progress → Send ALL tasks 2. Do the work immediately 3. Mark completed → Send ALL tasks -4. Wait for user approval +4. Wait for user approval of the completed task 5. If approved → Start next task (repeat from step 1) -6. If rejected → Redo the task +6. If rejected → Redo the task based on feedback Example (3 tasks total): Start task 1 - Send ALL: @@ -143,8 +143,9 @@ Start task 2 - Send ALL: Rules: - Send ALL tasks every single call (tool will reject partial lists) - Only ONE task "in_progress" at a time -- After approval, immediately start next task -- Continue autonomously until all tasks done`, +- After plan approval, start first task immediately +- Wait for approval after each task completion before starting next +- Continue autonomously through all tasks with approval checkpoints`, inputSchema: TaskWriteInputSchema, execute: async (input: TaskWriteInput): Promise => { try { @@ -188,16 +189,8 @@ Rules: approvalType = "plan"; // Create promise and emit approval request - const responsePromise = new Promise<{ approved: boolean; comment?: string }>((resolve, reject) => { + const responsePromise = new Promise<{ approved: boolean; comment?: string }>((resolve) => { pendingApprovalResolve = resolve; - - // Timeout after 5 minutes - setTimeout(() => { - if (pendingApprovalResolve) { - pendingApprovalResolve = null; - reject(new Error("Approval request timed out")); - } - }, 5 * 60 * 1000); }); eventHandler({ @@ -207,12 +200,7 @@ Rules: message: "Please review the implementation plan" }); - try { - approvalResult = await responsePromise; - } catch (error) { - console.error("[TaskWrite Tool] Approval failed:", error); - approvalResult = { approved: false, comment: "Approval request failed" }; - } + approvalResult = await responsePromise; } // Case 2: Task just completed (no in-progress tasks means agent finished work) // If there's an in_progress task, the agent is just starting work - no approval needed @@ -224,16 +212,8 @@ Rules: approvedTaskId = lastCompletedTask.id; // Create promise and emit approval request - const responsePromise = new Promise<{ approved: boolean; comment?: string }>((resolve, reject) => { + const responsePromise = new Promise<{ approved: boolean; comment?: string }>((resolve) => { pendingApprovalResolve = resolve; - - // Timeout after 5 minutes - setTimeout(() => { - if (pendingApprovalResolve) { - pendingApprovalResolve = null; - reject(new Error("Approval request timed out")); - } - }, 5 * 60 * 1000); }); eventHandler({ @@ -244,12 +224,7 @@ Rules: message: `Please verify the completed work for: ${lastCompletedTask.description}` }); - try { - approvalResult = await responsePromise; - } catch (error) { - console.error("[TaskWrite Tool] Approval failed:", error); - approvalResult = { approved: false, comment: "Approval request failed" }; - } + approvalResult = await responsePromise; } else if (inProgressTasks.length > 0) { console.log(`[TaskWrite Tool] Task in progress, no approval needed: ${inProgressTasks[0].description}`); } diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts index 1f21f451ef7..0f6f084ff8f 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/commandTemplates/data/commandTemplates.const.ts @@ -196,21 +196,21 @@ export const NATURAL_PROGRAMMING_TEMPLATES: TemplateDefinition[] = [ // Suggested command templates are defined here. export const suggestedCommandTemplates: AIPanelPrompt[] = [ { - type: 'command-template', + type: "command-template", command: Command.Code, templateId: TemplateId.Wildcard, - text: 'write a hello world http service', + text: "write a hello world http service", }, { - type: 'command-template', + type: "command-template", command: Command.Design, templateId: TemplateId.Wildcard, - text: 'design an API for a task management system', + text: "design an API for a task management system", }, { - type: 'command-template', + type: "command-template", command: Command.Ask, templateId: TemplateId.Wildcard, - text: 'how to write a concurrent application?', + text: "how to write a concurrent application?", }, ]; 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 86e72dc7be3..f8f1d00c042 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 @@ -328,10 +328,7 @@ const AIChat: React.FC = () => { // Show different indicator based on whether tasks exist if (!activeTasks || activeTasks.length === 0) { // Initial task creation - newMessages[newMessages.length - 1].content += `\n\nPlanning...`; - } else { - // Task update - newMessages[newMessages.length - 1].content += `\n\nUpdating tasks...`; + newMessages[newMessages.length - 1].content += ` Planning...`; } } return newMessages; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx index 239278f603c..4010596d382 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx @@ -17,7 +17,7 @@ */ import styled from "@emotion/styled"; -import React, { useState } from "react"; +import React, { useState, useEffect, useRef } from "react"; const DialogOverlay = styled.div` position: fixed; @@ -75,7 +75,7 @@ const TaskList = styled.div` const TaskItem = styled.div` display: flex; - align-items: flex-start; + align-items: center; gap: 10px; padding: 8px; margin-bottom: 8px; @@ -90,7 +90,7 @@ const TaskNumber = styled.span` flex-shrink: 0; `; -const TaskDescription = styled.div` +const TaskDescription = styled.span` flex: 1; color: var(--vscode-editor-foreground); line-height: 1.4; @@ -204,6 +204,9 @@ const ApprovalDialog: React.FC = ({ const [dialogState, setDialogState] = useState("initial"); const [comment, setComment] = useState(""); const [isSubmitting, setIsSubmitting] = useState(false); + const taskListRef = useRef(null); + const highlightedTaskRef = useRef(null); + const scrollTimeoutRef = useRef(null); const handleApprove = () => { setIsSubmitting(true); @@ -231,6 +234,50 @@ const ApprovalDialog: React.FC = ({ setComment(""); }; + // Function to scroll to highlighted task + const scrollToHighlightedTask = () => { + if (highlightedTaskRef.current) { + highlightedTaskRef.current.scrollIntoView({ + behavior: "smooth", + block: "nearest", + }); + } + }; + + // Auto-scroll to highlighted task on mount + useEffect(() => { + if (taskId && highlightedTaskRef.current) { + scrollToHighlightedTask(); + } + }, [taskId]); + + // Handle user scroll - refocus after delay + useEffect(() => { + const taskList = taskListRef.current; + if (!taskList || !taskId) return; + + const handleScroll = () => { + // Clear existing timeout + if (scrollTimeoutRef.current) { + clearTimeout(scrollTimeoutRef.current); + } + + // Set new timeout to refocus after 3 seconds + scrollTimeoutRef.current = setTimeout(() => { + scrollToHighlightedTask(); + }, 3000); + }; + + taskList.addEventListener("scroll", handleScroll); + + return () => { + taskList.removeEventListener("scroll", handleScroll); + if (scrollTimeoutRef.current) { + clearTimeout(scrollTimeoutRef.current); + } + }; + }, [taskId]); + const getDialogTitle = () => { if (approvalType === "plan") { return "Review Implementation Plan"; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx index a0a6dcc7956..f37391dae1b 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx @@ -18,7 +18,7 @@ import { keyframes } from "@emotion/css"; import styled from "@emotion/styled"; -import React, { useState } from "react"; +import React, { useState, useEffect, useRef } from "react"; const spin = keyframes` from { transform: rotate(0deg); } @@ -32,24 +32,24 @@ const TodoContainer = styled.div` padding: 0; margin: 0; font-family: var(--vscode-editor-font-family); - font-size: 13px; + font-size: 12px; color: var(--vscode-editor-foreground); - min-height: 32px; + min-height: 24px; `; const TodoHeader = styled.div<{ clickable?: boolean }>` font-weight: 600; - margin-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '12px'}; - padding-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '8px'}; + margin-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '6px'}; + padding-bottom: ${(props: { clickable?: boolean }) => props.clickable ? '0' : '4px'}; border-bottom: ${(props: { clickable?: boolean }) => props.clickable ? 'none' : '1px solid var(--vscode-panel-border)'}; display: flex; align-items: center; - gap: 8px; + gap: 4px; cursor: ${(props: { clickable?: boolean }) => props.clickable ? 'pointer' : 'default'}; user-select: none; - padding: 4px; - border-radius: 4px; + padding: 2px; + border-radius: 3px; &:hover { background-color: ${(props: { clickable?: boolean }) => @@ -67,22 +67,24 @@ const ChevronIcon = styled.span<{ expanded: boolean }>` const MinimalTaskInfo = styled.span` color: var(--vscode-descriptionForeground); font-weight: 400; - font-size: 12px; + font-size: 11px; margin-left: 4px; `; const TodoList = styled.div` display: flex; flex-direction: column; - gap: 8px; + gap: 4px; + overflow-y: auto; + max-height: 200px; `; const TodoItem = styled.div<{ status: string }>` display: flex; - align-items: flex-start; - gap: 10px; - padding: 8px; - border-radius: 4px; + align-items: center; + gap: 6px; + padding: 4px 6px; + border-radius: 3px; background-color: ${(props: { status: string }) => props.status === "completed" ? "var(--vscode-list-hoverBackground)" @@ -96,9 +98,8 @@ const TodoIcon = styled.span<{ status: string }>` display: flex; align-items: center; justify-content: center; - width: 18px; - height: 18px; - margin-top: 2px; + width: 16px; + height: 16px; &.pending { .codicon { @@ -120,16 +121,17 @@ const TodoIcon = styled.span<{ status: string }>` } `; -const TodoText = styled.div<{ status: string }>` +const TodoText = styled.span<{ status: string }>` flex: 1; text-decoration: ${(props: { status: string }) => props.status === "completed" ? "line-through" : "none"}; + line-height: 16px; `; const TodoNumber = styled.span` color: var(--vscode-descriptionForeground); font-weight: 600; - margin-right: 4px; + margin-right: 2px; `; export interface Task { @@ -157,6 +159,9 @@ const getStatusIcon = (status: string): { className: string; icon: string } => { const TodoSection: React.FC = ({ tasks, message }) => { const [isExpanded, setIsExpanded] = useState(true); + const inProgressRef = useRef(null); + const todoListRef = useRef(null); + const scrollTimeoutRef = useRef(null); const completedCount = tasks.filter((t) => t.status === "completed").length; const inProgressTask = tasks.find((t) => t.status === "in_progress"); const allCompleted = completedCount === tasks.length; @@ -166,6 +171,50 @@ const TodoSection: React.FC = ({ tasks, message }) => { setIsExpanded(!isExpanded); }; + // Function to scroll to in-progress task + const scrollToInProgress = () => { + if (inProgressRef.current) { + inProgressRef.current.scrollIntoView({ + behavior: "smooth", + block: "nearest", + }); + } + }; + + // Auto-scroll to in-progress task + useEffect(() => { + if (isExpanded && hasInProgress) { + scrollToInProgress(); + } + }, [isExpanded, inProgressTask?.id]); + + // Handle user scroll - refocus after delay + useEffect(() => { + const todoList = todoListRef.current; + if (!todoList || !hasInProgress) return; + + const handleScroll = () => { + // Clear existing timeout + if (scrollTimeoutRef.current) { + clearTimeout(scrollTimeoutRef.current); + } + + // Set new timeout to refocus after 3 seconds + scrollTimeoutRef.current = setTimeout(() => { + scrollToInProgress(); + }, 3000); + }; + + todoList.addEventListener("scroll", handleScroll); + + return () => { + todoList.removeEventListener("scroll", handleScroll); + if (scrollTimeoutRef.current) { + clearTimeout(scrollTimeoutRef.current); + } + }; + }, [hasInProgress, inProgressTask?.id]); + // Determine status text const getStatusText = () => { if (allCompleted) return "completed"; @@ -195,9 +244,9 @@ const TodoSection: React.FC = ({ tasks, message }) => { {message && (
= ({ tasks, message }) => { {message}
)} - + {tasks.map((task, index) => { const statusInfo = getStatusIcon(task.status); + const isInProgress = task.status === "in_progress"; return ( - + From d0192d2652b0de2f2c265ab33a56335b40ee742b Mon Sep 17 00:00:00 2001 From: RNViththagan Date: Tue, 21 Oct 2025 20:33:46 +0530 Subject: [PATCH 006/619] Enables adding code to workspace from approval dialog Adds functionality to allow users to add code segments directly to the workspace from the AI completion approval dialog. This change introduces a button in the approval dialog that, when clicked, adds the code segments associated with the AI's completion to the user's workspace. It extracts code segments from the last assistant message and adds them to the workspace using existing handler functions. The button is only displayed for "completion" approval types and is disabled after the code has been successfully added or if an error occurs during the process. --- .../views/AIPanel/components/AIChat/index.tsx | 117 ++++++++++++++++-- .../AIPanel/components/ApprovalDialog.tsx | 99 ++++++++++----- .../views/AIPanel/components/TodoSection.tsx | 24 +++- .../src/views/AIPanel/styles.tsx | 4 +- 4 files changed, 201 insertions(+), 43 deletions(-) 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 f8f1d00c042..d5f989ac600 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 @@ -161,6 +161,7 @@ const AIChat: React.FC = () => { tasks: any[]; taskId?: string; message?: string; + codeSegments?: any[]; } | null>(null); // Top-level tasks state for the todo panel @@ -445,11 +446,53 @@ const AIChat: React.FC = () => { } else if (type === "task_approval_request") { // Handle approval request from the backend console.log("[Approval] Received approval request:", response); + + // Only extract code segments for completion approval (not plan approval) + const codeSegments: any[] = []; + + if (response.approvalType === "completion") { + setMessages((currentMessages) => { + console.log("[Approval] Total current messages:", currentMessages.length); + + const lastAssistantMessage = currentMessages + .slice() + .reverse() + .find((msg: any) => msg.role === "Copilot"); + + console.log("[Approval] Last Copilot message:", lastAssistantMessage ? "found" : "not found"); + + if (lastAssistantMessage) { + const segments = splitContent(lastAssistantMessage.content); + console.log("[Approval] Found total segments:", segments.length); + + // Efficiently filter and map only Code segments using filter + map + const extractedCodeSegments = segments + .filter((segment) => segment.type === SegmentType.Code) + .map((segment) => ({ + segmentText: segment.text.trim(), + filePath: segment.fileName, + language: segment.language, + })); + + codeSegments.push(...extractedCodeSegments); + console.log("[Approval] Extracted code segments:", codeSegments.length); + } + + // Return unchanged messages + return currentMessages; + }); + } + + console.log("[Approval] Approval type:", response.approvalType); + console.log("[Approval] Final code segments count:", codeSegments.length); + + // Set approval request after extracting code segments setApprovalRequest({ approvalType: response.approvalType, tasks: response.tasks, taskId: response.taskId, message: response.message, + codeSegments: codeSegments.length > 0 ? codeSegments : undefined, }); } else if (type === "intermediary_state") { const state = response.state; @@ -2079,6 +2122,46 @@ const AIChat: React.FC = () => { setApprovalRequest(null); }; + // Handle adding code to workspace from approval dialog + const handleAddCodeToWorkspaceFromApproval = async () => { + if (!approvalRequest || !approvalRequest.codeSegments || approvalRequest.codeSegments.length === 0) { + console.log("[Approval] No code segments to add"); + return; + } + + console.log("[Approval] Adding code segments to workspace:", approvalRequest.codeSegments); + + // Get the command from the latest message + const lastAssistantMessage = messagesRef.current + .slice() + .reverse() + .find((msg: any) => msg.role === "assistant"); + + let command = "ai_design"; // default command + if (lastAssistantMessage) { + const segments = splitContent(lastAssistantMessage.content); + // Find the first code segment that has a command + const codeSegmentWithCommand = segments.find((s) => s.type === SegmentType.Code && s.command); + if (codeSegmentWithCommand && codeSegmentWithCommand.command) { + command = codeSegmentWithCommand.command; + } + } + + // Create a dummy setState function since we don't need to track individual code section state + const dummySetIsCodeAdded = () => {}; + + try { + await handleAddAllCodeSegmentsToWorkspace( + approvalRequest.codeSegments, + dummySetIsCodeAdded, + command + ); + console.log("[Approval] Successfully added code segments to workspace"); + } catch (error) { + console.error("[Approval] Error adding code segments to workspace:", error); + } + }; + return ( <> {!showSettings && ( @@ -2377,16 +2460,30 @@ const AIChat: React.FC = () => { )} {showSettings && setShowSettings(false)}>} - {approvalRequest && ( - - )} + {approvalRequest && (() => { + const shouldShowAddButton = approvalRequest.approvalType === "completion" && approvalRequest.codeSegments; + console.log("[ApprovalDialog Render]", { + approvalType: approvalRequest.approvalType, + hasCodeSegments: !!approvalRequest.codeSegments, + codeSegmentsLength: approvalRequest.codeSegments?.length, + shouldShowAddButton, + }); + return ( + + ); + })()} ); }; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx index 4010596d382..35a2d40d331 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx @@ -189,6 +189,7 @@ interface ApprovalDialogProps { message?: string; onApprove: (comment?: string) => void; onReject: (comment?: string) => void; + onAddToWorkspace?: () => void; } type DialogState = "initial" | "rejecting"; @@ -200,14 +201,23 @@ const ApprovalDialog: React.FC = ({ message, onApprove, onReject, + onAddToWorkspace, }) => { const [dialogState, setDialogState] = useState("initial"); const [comment, setComment] = useState(""); const [isSubmitting, setIsSubmitting] = useState(false); + const [isAddingToWorkspace, setIsAddingToWorkspace] = useState(false); + const [hasAddedToWorkspace, setHasAddedToWorkspace] = useState(false); const taskListRef = useRef(null); const highlightedTaskRef = useRef(null); const scrollTimeoutRef = useRef(null); + console.log("[ApprovalDialog] Props:", { + approvalType, + hasAddToWorkspace: !!onAddToWorkspace, + taskCount: tasks.length, + }); + const handleApprove = () => { setIsSubmitting(true); onApprove(undefined); // No comment when approving @@ -234,6 +244,21 @@ const ApprovalDialog: React.FC = ({ setComment(""); }; + const handleAddToWorkspace = async () => { + if (onAddToWorkspace && !hasAddedToWorkspace) { + setIsAddingToWorkspace(true); + try { + await onAddToWorkspace(); + setHasAddedToWorkspace(true); // Disable button after successful addition + } catch (error) { + console.error("[ApprovalDialog] Error adding to workspace:", error); + // Don't set hasAddedToWorkspace on error, so user can retry + } finally { + setIsAddingToWorkspace(false); + } + } + }; + // Function to scroll to highlighted task const scrollToHighlightedTask = () => { if (highlightedTaskRef.current) { @@ -313,37 +338,39 @@ const ApprovalDialog: React.FC = ({ {getDialogMessage()} - - {tasks.map((task, index) => ( - + {tasks.map((task, index) => { + const isHighlighted = task.id === taskId; + return ( + - {index + 1}. - - {task.description} - {task.status === "completed" && ( - - - - )} - - - ))} + }} + > + {index + 1}. + + {task.description} + {task.status === "completed" && ( + + + + )} + + + ); + })} {dialogState === "rejecting" && ( @@ -364,10 +391,24 @@ const ApprovalDialog: React.FC = ({ {dialogState === "initial" ? ( <> + {approvalType === "completion" && onAddToWorkspace && ( + + )} - {activeTasks && activeTasks.length > 0 && ( + {/* Show TodoSection when we have tasks from approval request OR after plan is approved */} + {(approvalRequest?.tasks || todoTasks) && ( - + )}
@@ -2482,30 +2448,6 @@ const AIChat: React.FC = () => { )} {showSettings && setShowSettings(false)}>} - {approvalRequest && (() => { - const shouldShowAddButton = approvalRequest.approvalType === "completion" && approvalRequest.codeSegments; - console.log("[ApprovalDialog Render]", { - approvalType: approvalRequest.approvalType, - hasCodeSegments: !!approvalRequest.codeSegments, - codeSegmentsLength: approvalRequest.codeSegments?.length, - shouldShowAddButton, - }); - return ( - - ); - })()} ); }; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx deleted file mode 100644 index 9f8f480b7c7..00000000000 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/ApprovalDialog.tsx +++ /dev/null @@ -1,471 +0,0 @@ -/** - * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com) All Rights Reserved. - * - * WSO2 LLC. licenses this file to you under the Apache License, - * Version 2.0 (the "License"); you may not use this file except - * in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -import styled from "@emotion/styled"; -import React, { useState, useEffect, useRef } from "react"; - -const DialogOverlay = styled.div` - position: fixed; - top: 0; - left: 0; - right: 0; - bottom: 0; - background-color: rgba(0, 0, 0, 0.5); - display: flex; - align-items: center; - justify-content: center; - z-index: 1000; -`; - -const DialogContainer = styled.div` - background-color: var(--vscode-editor-background); - border: 1px solid var(--vscode-panel-border); - border-radius: 8px; - padding: 20px; - max-width: 600px; - width: 90%; - max-height: 80vh; - overflow-y: auto; - box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3); -`; - -const DialogHeader = styled.div` - display: flex; - align-items: center; - gap: 10px; - margin-bottom: 16px; - padding-bottom: 12px; - border-bottom: 1px solid var(--vscode-panel-border); -`; - -const DialogTitle = styled.h3` - margin: 0; - font-size: 16px; - font-weight: 600; - color: var(--vscode-editor-foreground); -`; - -const DialogMessage = styled.div` - margin-bottom: 16px; - font-size: 13px; - color: var(--vscode-descriptionForeground); - line-height: 1.5; -`; - -const TaskList = styled.div` - margin-bottom: 16px; - max-height: 300px; - overflow-y: auto; -`; - -const TaskItem = styled.div` - display: flex; - align-items: center; - gap: 10px; - padding: 8px; - margin-bottom: 8px; - background-color: var(--vscode-textCodeBlock-background); - border-radius: 4px; - font-size: 13px; -`; - -const TaskNumber = styled.span` - color: var(--vscode-descriptionForeground); - font-weight: 600; - flex-shrink: 0; -`; - -const TaskDescription = styled.span` - flex: 1; - color: var(--vscode-editor-foreground); - line-height: 1.4; -`; - -const CommentSection = styled.div` - margin-bottom: 16px; -`; - -const CommentLabel = styled.label` - display: block; - margin-bottom: 8px; - font-size: 13px; - color: var(--vscode-editor-foreground); - font-weight: 500; -`; - -const CommentTextarea = styled.textarea` - width: 100%; - min-height: 80px; - padding: 8px; - background-color: var(--vscode-input-background); - color: var(--vscode-input-foreground); - border: 1px solid var(--vscode-input-border); - border-radius: 4px; - font-family: var(--vscode-editor-font-family); - font-size: 13px; - resize: vertical; - box-sizing: border-box; - - &:focus { - outline: none; - border-color: var(--vscode-focusBorder); - } - - &::placeholder { - color: var(--vscode-input-placeholderForeground); - } -`; - -const ButtonGroup = styled.div` - display: flex; - gap: 12px; - justify-content: flex-end; -`; - -const Button = styled.button<{ variant?: "primary" | "secondary" }>` - padding: 8px 16px; - font-size: 13px; - font-weight: 500; - border: none; - border-radius: 4px; - cursor: pointer; - transition: all 0.2s ease; - font-family: var(--vscode-editor-font-family); - - ${(props: { variant?: "primary" | "secondary" }) => - props.variant === "primary" - ? ` - background-color: var(--vscode-button-background); - color: var(--vscode-button-foreground); - - &:hover { - background-color: var(--vscode-button-hoverBackground); - } - ` - : ` - background-color: var(--vscode-button-secondaryBackground); - color: var(--vscode-button-secondaryForeground); - - &:hover { - background-color: var(--vscode-button-secondaryHoverBackground); - } - `} - - &:active { - opacity: 0.8; - } - - &:disabled { - opacity: 0.5; - cursor: not-allowed; - } -`; - -export interface Task { - id: string; - description: string; - status: "pending" | "in_progress" | "review" | "done" | "rejected"; -} - -interface ApprovalDialogProps { - approvalType: "plan" | "completion"; - tasks: Task[]; - taskId?: string; - message?: string; - onApprove: (comment?: string) => void; - onReject: (comment?: string) => void; - onAddToWorkspace?: () => void; -} - -type DialogState = "initial" | "rejecting"; - -const ApprovalDialog: React.FC = ({ - approvalType, - tasks, - taskId, - message, - onApprove, - onReject, - onAddToWorkspace, -}) => { - const [dialogState, setDialogState] = useState("initial"); - const [comment, setComment] = useState(""); - const [isSubmitting, setIsSubmitting] = useState(false); - const [isAddingToWorkspace, setIsAddingToWorkspace] = useState(false); - const [hasAddedToWorkspace, setHasAddedToWorkspace] = useState(false); - const taskListRef = useRef(null); - const highlightedTaskRef = useRef(null); - const scrollTimeoutRef = useRef(null); - - console.log("[ApprovalDialog] Props:", { - approvalType, - hasAddToWorkspace: !!onAddToWorkspace, - taskCount: tasks.length, - }); - - const handleApprove = () => { - setIsSubmitting(true); - onApprove(undefined); // No comment when approving - }; - - const handleRejectClick = () => { - // Switch to reject state to show comment field - setDialogState("rejecting"); - }; - - const handleRejectSubmit = () => { - const trimmedComment = comment.trim(); - if (!trimmedComment) { - // Comment is required for rejection - return; - } - setIsSubmitting(true); - onReject(trimmedComment); - }; - - const handleBack = () => { - // Go back to initial state from rejecting state - setDialogState("initial"); - setComment(""); - }; - - const handleAddToWorkspace = async () => { - if (onAddToWorkspace && !hasAddedToWorkspace) { - setIsAddingToWorkspace(true); - try { - await onAddToWorkspace(); - setHasAddedToWorkspace(true); // Disable button after successful addition - } catch (error) { - console.error("[ApprovalDialog] Error adding to workspace:", error); - // Don't set hasAddedToWorkspace on error, so user can retry - } finally { - setIsAddingToWorkspace(false); - } - } - }; - - // Function to scroll to highlighted task - const scrollToHighlightedTask = () => { - if (highlightedTaskRef.current) { - highlightedTaskRef.current.scrollIntoView({ - behavior: "smooth", - block: "nearest", - }); - } - }; - - // Auto-scroll to highlighted task on mount - useEffect(() => { - if (taskId && highlightedTaskRef.current) { - scrollToHighlightedTask(); - } - }, [taskId]); - - // Handle user scroll - refocus after delay - useEffect(() => { - const taskList = taskListRef.current; - if (!taskList || !taskId) return; - - const handleScroll = () => { - // Clear existing timeout - if (scrollTimeoutRef.current) { - clearTimeout(scrollTimeoutRef.current); - } - - // Set new timeout to refocus after 3 seconds - scrollTimeoutRef.current = setTimeout(() => { - scrollToHighlightedTask(); - }, 3000); - }; - - taskList.addEventListener("scroll", handleScroll); - - return () => { - taskList.removeEventListener("scroll", handleScroll); - if (scrollTimeoutRef.current) { - clearTimeout(scrollTimeoutRef.current); - } - }; - }, [taskId]); - - const getDialogTitle = () => { - if (approvalType === "plan") { - return "Review Implementation Plan"; - } else { - const task = tasks.find((t) => t.id === taskId); - return `Review Completed Task: ${task?.description || "Unknown Task"}`; - } - }; - - const getDialogMessage = () => { - if (message) { - return message; - } - if (approvalType === "plan") { - return "Please review the proposed implementation plan below. You can approve to proceed or reject with feedback for revisions."; - } else { - return "Please review the completed work. You can approve to continue or reject with feedback for corrections."; - } - }; - - return ( - - - - - {getDialogTitle()} - - - {getDialogMessage()} - - - {tasks.map((task, index) => { - const isHighlighted = task.id === taskId; - return ( - - {index + 1}. - - {task.description} - {task.status === "done" && ( - - - - )} - {task.status === "review" && ( - - - - )} - {task.status === "in_progress" && ( - - - - )} - - - ); - })} - - - {dialogState === "rejecting" && ( - - - Please provide feedback for rejection (required): - - setComment(e.target.value)} - placeholder="Enter the reason for rejection and what needs to be changed..." - disabled={isSubmitting} - autoFocus - /> - - )} - - - {dialogState === "initial" ? ( - <> - {approvalType === "completion" && onAddToWorkspace && ( - - )} - - - - ) : ( - <> - - - - )} - - - - ); -}; - -export default ApprovalDialog; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx index 3e0e8bede25..b5a39d06617 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/TodoSection.tsx @@ -25,17 +25,9 @@ const spin = keyframes` to { transform: rotate(360deg); } `; -const flashHighlight = keyframes` - 0% { - background-color: rgba(var(--vscode-list-hoverBackground-rgb, 128, 128, 128), 0.3); - } - 100% { - background-color: transparent; - } -`; - const TodoContainer = styled.div<{ isNew?: boolean }>` - background-color: transparent; + background-color: ${(props: { isNew?: boolean }) => + props.isNew ? 'rgba(128, 128, 128, 0.3)' : 'transparent'}; border: none; border-radius: 0; padding: 0; @@ -44,8 +36,7 @@ const TodoContainer = styled.div<{ isNew?: boolean }>` font-size: 12px; color: var(--vscode-editor-foreground); min-height: 24px; - animation: ${(props: { isNew?: boolean }) => - props.isNew ? `${flashHighlight} 3s ease-out 1` : 'none'}; + transition: background-color 0.3s ease-out; `; const TodoHeader = styled.div<{ clickable?: boolean }>` @@ -145,6 +136,92 @@ const TodoNumber = styled.span` margin-right: 2px; `; +const ApprovalSection = styled.div` + margin-top: 8px; + padding-top: 8px; + border-top: 1px solid var(--vscode-panel-border); +`; + +const ButtonGroup = styled.div` + display: flex; + gap: 8px; + justify-content: flex-end; +`; + +const Button = styled.button<{ variant?: "primary" | "secondary" }>` + padding: 6px 12px; + font-size: 12px; + font-weight: 500; + border: none; + border-radius: 3px; + cursor: pointer; + transition: all 0.2s ease; + font-family: var(--vscode-editor-font-family); + display: flex; + align-items: center; + gap: 4px; + + ${(props: { variant?: "primary" | "secondary" }) => + props.variant === "primary" + ? ` + background-color: var(--vscode-button-background); + color: var(--vscode-button-foreground); + + &:hover { + background-color: var(--vscode-button-hoverBackground); + } + ` + : ` + background-color: var(--vscode-button-secondaryBackground); + color: var(--vscode-button-secondaryForeground); + + &:hover { + background-color: var(--vscode-button-secondaryHoverBackground); + } + `} + + &:active { + opacity: 0.8; + } + + &:disabled { + opacity: 0.5; + cursor: not-allowed; + } +`; + +const CommentTextarea = styled.textarea` + width: 100%; + min-height: 60px; + padding: 6px; + margin-bottom: 8px; + background-color: var(--vscode-input-background); + color: var(--vscode-input-foreground); + border: 1px solid var(--vscode-input-border); + border-radius: 3px; + font-family: var(--vscode-editor-font-family); + font-size: 12px; + resize: vertical; + box-sizing: border-box; + + &:focus { + outline: none; + border-color: var(--vscode-focusBorder); + } + + &::placeholder { + color: var(--vscode-input-placeholderForeground); + } +`; + +const CommentLabel = styled.label` + display: block; + margin-bottom: 6px; + font-size: 11px; + color: var(--vscode-editor-foreground); + font-weight: 500; +`; + export interface Task { id: string; description: string; @@ -154,6 +231,9 @@ export interface Task { interface TodoSectionProps { tasks: Task[]; message?: string; + onApprove?: (comment?: string) => void; + onReject?: (comment?: string) => void; + approvalType?: "plan" | "completion"; } const getStatusIcon = (status: string): { className: string; icon: string } => { @@ -172,25 +252,36 @@ const getStatusIcon = (status: string): { className: string; icon: string } => { } }; -const TodoSection: React.FC = ({ tasks, message }) => { +const TodoSection: React.FC = ({ tasks, message, onApprove, onReject, approvalType }) => { const [isExpanded, setIsExpanded] = useState(true); - const [isNew, setIsNew] = useState(true); + const [isNew, setIsNew] = useState(false); + const [showRejectComment, setShowRejectComment] = useState(false); + const [rejectComment, setRejectComment] = useState(""); + const [isSubmitting, setIsSubmitting] = useState(false); const inProgressRef = useRef(null); const todoListRef = useRef(null); const scrollTimeoutRef = useRef(null); const completedCount = tasks.filter((t) => t.status === "done").length; const inProgressTask = tasks.find((t) => t.status === "in_progress"); - const reviewTask = tasks.find((t) => t.status === "review"); const allCompleted = completedCount === tasks.length; const hasInProgress = !!inProgressTask; + const needsApproval = !!approvalType; - // Remove the "new" state after animation completes + // Highlight container only for plan approval, not task completion approval useEffect(() => { - const timer = setTimeout(() => { + if (approvalType === "plan") { + // Highlight entire container for plan approval + setIsNew(true); + } else { + // No container highlight for task approval (task itself has blue border) setIsNew(false); - }, 3000); // 1.5s * 2 iterations = 3s - return () => clearTimeout(timer); - }, []); + } + + // Reset submission state when approval type changes + setIsSubmitting(false); + setShowRejectComment(false); + setRejectComment(""); + }, [approvalType]); const toggleExpanded = () => { setIsExpanded(!isExpanded); @@ -243,10 +334,36 @@ const TodoSection: React.FC = ({ tasks, message }) => { // Determine status text const getStatusText = () => { if (allCompleted) return "completed"; + if (needsApproval) return "awaiting approval"; if (hasInProgress) return "in progress"; return "ongoing"; }; + const handleApprove = () => { + if (onApprove) { + setIsSubmitting(true); + onApprove(undefined); + } + }; + + const handleRejectClick = () => { + setShowRejectComment(true); + }; + + const handleRejectSubmit = () => { + const trimmedComment = rejectComment.trim(); + if (!trimmedComment || !onReject) { + return; + } + setIsSubmitting(true); + onReject(trimmedComment); + }; + + const handleRejectCancel = () => { + setShowRejectComment(false); + setRejectComment(""); + }; + return ( @@ -283,11 +400,20 @@ const TodoSection: React.FC = ({ tasks, message }) => { {tasks.map((task, index) => { const statusInfo = getStatusIcon(task.status); const isInProgress = task.status === "in_progress"; + const isReview = task.status === "review"; return ( @@ -302,6 +428,61 @@ const TodoSection: React.FC = ({ tasks, message }) => { )} + {needsApproval && ( + + {showRejectComment ? ( + <> + + Provide feedback for revision (required): + + setRejectComment(e.target.value)} + placeholder="Enter what needs to be changed..." + disabled={isSubmitting} + autoFocus + /> + + + + + + ) : ( + + + + + )} + + )} ); }; diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx index 6b67e3a6ce3..e05cab3ce56 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/styles.tsx @@ -46,12 +46,12 @@ export const HeaderButtons = styled.div({ }); export const TodoPanel = styled.div` - max-height: 30vh; - overflow-y: auto; border-bottom: 1px solid rgba(128, 128, 128, 0.3); background-color: var(--vscode-editor-background); padding: 8px 12px; flex-shrink: 0; + display: flex; + flex-direction: column; `; export const Main = styled.main({ From 7476be56c1f8c601fb18e505d374afabbb315f2f Mon Sep 17 00:00:00 2001 From: RNViththagan Date: Sat, 25 Oct 2025 13:35:03 +0530 Subject: [PATCH 010/619] Add auto approval for tasks in design mode with UI toggle Introduces automatic task approval functionality for design mode with user control toggle. - Add ENABLE_AUTO_APPROVE and DISABLE_AUTO_APPROVE events to state machine - Include autoApproveEnabled flag in AI chat machine context - Implement auto approval logic in task write tool to bypass manual approval - Add UI toggle button in AI chat interface to control auto approval mode - Enhance RPC layer with context retrieval method for state synchronization - Update task management workflow to support both automatic and manual approval flows - Display auto approval status in UI with clear on/off indication --- .../ballerina-core/src/state-machine-types.ts | 6 + .../ballerina-extension/src/RPCLayer.ts | 3 +- .../ai/service/libs/task_write_tool.ts | 117 ++++++++++++------ .../src/views/ai-panel/aiChatMachine.ts | 17 ++- .../src/BallerinaRpcClient.ts | 8 +- .../views/AIPanel/components/AIChat/index.tsx | 39 ++++++ 6 files changed, 146 insertions(+), 44 deletions(-) diff --git a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts index 4541332cd18..a50c89b6124 100644 --- a/workspaces/ballerina/ballerina-core/src/state-machine-types.ts +++ b/workspaces/ballerina/ballerina-core/src/state-machine-types.ts @@ -359,6 +359,8 @@ export enum AIChatMachineEventType { TASK_COMPLETED = 'TASK_COMPLETED', FINISH_EXECUTION = 'FINISH_EXECUTION', APPROVE_TASK = 'APPROVE_TASK', + ENABLE_AUTO_APPROVE = 'ENABLE_AUTO_APPROVE', + DISABLE_AUTO_APPROVE = 'DISABLE_AUTO_APPROVE', REJECT_TASK = 'REJECT_TASK', RESTORE_STATE = 'RESTORE_STATE', ERROR = 'ERROR', @@ -427,6 +429,7 @@ export interface AIChatMachineContext { sessionId?: string; projectId?: string; currentApproval?: UserApproval; + autoApproveEnabled?: boolean; } export type AIChatMachineSendableEvent = @@ -439,6 +442,8 @@ export type AIChatMachineSendableEvent = | { type: AIChatMachineEventType.TASK_COMPLETED } | { type: AIChatMachineEventType.FINISH_EXECUTION } | { type: AIChatMachineEventType.APPROVE_TASK; payload?: { comment?: string; lastApprovedTaskIndex?: number } } + | { type: AIChatMachineEventType.ENABLE_AUTO_APPROVE } + | { type: AIChatMachineEventType.DISABLE_AUTO_APPROVE } | { type: AIChatMachineEventType.REJECT_TASK; payload: { comment?: string } } | { type: AIChatMachineEventType.RESET } | { type: AIChatMachineEventType.RESTORE_STATE; payload: { state: AIChatMachineContext } } @@ -504,3 +509,4 @@ export const currentThemeChanged: NotificationType = { method: ' export const aiChatStateChanged: NotificationType = { method: 'aiChatStateChanged' }; export const sendAIChatStateEvent: RequestType = { method: 'sendAIChatStateEvent' }; +export const getAIChatContext: RequestType = { method: 'getAIChatContext' }; diff --git a/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts b/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts index b074dda074d..f05baf8440a 100644 --- a/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts +++ b/workspaces/ballerina/ballerina-extension/src/RPCLayer.ts @@ -19,7 +19,7 @@ import { WebviewView, WebviewPanel, window } from 'vscode'; import { Messenger } from 'vscode-messenger'; import { StateMachine } from './stateMachine'; -import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent, aiChatStateChanged, sendAIChatStateEvent, AIChatMachineEventType, AIChatMachineSendableEvent } from '@wso2/ballerina-core'; +import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent, aiChatStateChanged, sendAIChatStateEvent, getAIChatContext, AIChatMachineEventType, AIChatMachineSendableEvent } from '@wso2/ballerina-core'; import { VisualizerWebview } from './views/visualizer/webview'; import { registerVisualizerRpcHandlers } from './rpc-managers/visualizer/rpc-handler'; import { registerLangClientRpcHandlers } from './rpc-managers/lang-client/rpc-handler'; @@ -101,6 +101,7 @@ export class RPCLayer { registerAiPanelRpcHandlers(RPCLayer._messenger); RPCLayer._messenger.onRequest(sendAIStateEvent, (event: AIMachineEventType | AIMachineSendableEvent) => AIStateMachine.sendEvent(event)); RPCLayer._messenger.onRequest(sendAIChatStateEvent, (event: AIChatMachineEventType | AIChatMachineSendableEvent) => AIChatStateMachine.sendEvent(event)); + RPCLayer._messenger.onRequest(getAIChatContext, () => AIChatStateMachine.context()); // ----- Data Mapper Webview RPC Methods registerDataMapperRpcHandlers(RPCLayer._messenger); 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 40cecbc7a75..668be93206e 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 @@ -302,55 +302,90 @@ Rules: type: AIChatMachineEventType.TASK_COMPLETED, }); - // Also emit UI event for task approval display - eventHandler({ - type: "task_approval_request", - approvalType: "completion", - tasks: allTasks, - taskId: lastReviewTask.id, - message: `Please verify the completed work for: ${lastReviewTask.description}` - }); - - // Wait for user approval - state machine will transition out of TaskReview - console.log("[TaskWrite Tool] Waiting for task approval/rejection..."); + // Check if auto-approval is enabled + const isAutoApproveEnabled = currentContext.autoApproveEnabled === true; + + if (isAutoApproveEnabled) { + // Auto-approval mode: automatically approve the task + console.log( + `[TaskWrite Tool] Auto-approval enabled, auto-approving task: ${lastReviewTask.id}` + ); + + // Immediately emit APPROVE_TASK event to auto-approve without waiting for user + AIChatStateMachine.sendEvent({ + type: AIChatMachineEventType.APPROVE_TASK, + }); - approvalResult = await new Promise<{ approved: boolean; comment?: string }>((resolve) => { - // Subscribe to state machine changes - const subscription = AIChatStateMachine.service().subscribe((state) => { - const currentState = state.value; + // Update all review tasks to done status immediately + reviewTasks.forEach((task) => { + const taskIndex = allTasks.findIndex((t) => t.id === task.id); + if (taskIndex !== -1) { + allTasks[taskIndex].status = TaskStatus.DONE; + } + }); - // If we moved from TaskReview to ApprovedTask = approved - if (currentState === 'ApprovedTask') { - console.log(`[TaskWrite Tool] Task(s) approved by user (${reviewTasks.length} task(s))`); + // Set approval result without waiting + approvalResult = { approved: true }; + } else { + // Manual approval mode: emit event and wait for user approval + console.log(`[TaskWrite Tool] Manual approval mode, waiting for user approval`); + + // Also emit UI event for task approval display + eventHandler({ + type: "task_approval_request", + approvalType: "completion", + tasks: allTasks, + taskId: lastReviewTask.id, + message: `Please verify the completed work for: ${lastReviewTask.description}`, + }); - // Update all review tasks to done status - reviewTasks.forEach(task => { - const taskIndex = allTasks.findIndex(t => t.id === task.id); + // Wait for user approval - state machine will transition out of TaskReview + console.log("[TaskWrite Tool] Waiting for task approval/rejection..."); + + approvalResult = await new Promise<{ approved: boolean; comment?: string }>((resolve) => { + // Subscribe to state machine changes + const subscription = AIChatStateMachine.service().subscribe((state) => { + const currentState = state.value; + + // If we moved from TaskReview to ApprovedTask = approved + if (currentState === "ApprovedTask") { + console.log( + `[TaskWrite Tool] Task(s) approved by user (${reviewTasks.length} task(s))` + ); + + // Update all review tasks to done status + reviewTasks.forEach((task) => { + const taskIndex = allTasks.findIndex((t) => t.id === task.id); + if (taskIndex !== -1) { + allTasks[taskIndex].status = TaskStatus.DONE; + } + }); + + subscription.unsubscribe(); + resolve({ approved: true }); + } + // If we moved from TaskReview to RejectedTask = rejected + else if (currentState === "RejectedTask") { + // Get rejection comment from state machine context + const rejectionComment = state.context.currentApproval?.comment; + console.log( + `[TaskWrite Tool] Task rejected by user${ + rejectionComment ? `: "${rejectionComment}"` : "" + }` + ); + + // Update the last review task (the one being rejected) to rejected status + const taskIndex = allTasks.findIndex((t) => t.id === lastReviewTask.id); if (taskIndex !== -1) { - allTasks[taskIndex].status = TaskStatus.DONE; + allTasks[taskIndex].status = TaskStatus.REJECTED; } - }); - - subscription.unsubscribe(); - resolve({ approved: true }); - } - // If we moved from TaskReview to RejectedTask = rejected - else if (currentState === 'RejectedTask') { - // Get rejection comment from state machine context - const rejectionComment = state.context.currentApproval?.comment; - console.log(`[TaskWrite Tool] Task rejected by user${rejectionComment ? `: "${rejectionComment}"` : ''}`); - // Update the last review task (the one being rejected) to rejected status - const taskIndex = allTasks.findIndex(t => t.id === lastReviewTask.id); - if (taskIndex !== -1) { - allTasks[taskIndex].status = TaskStatus.REJECTED; + subscription.unsubscribe(); + resolve({ approved: false, comment: rejectionComment }); } - - subscription.unsubscribe(); - resolve({ approved: false, comment: rejectionComment }); - } + }); }); - }); + } } else if (inProgressTasks.length > 0) { // Emit START_TASK_EXECUTION event to state machine (non-blocking) AIChatStateMachine.sendEvent({ diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts index f837586c67d..3f7f39df7f2 100644 --- a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts +++ b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts @@ -83,8 +83,23 @@ const chatMachine = createMachine true, + chatHistory: (ctx) => + addChatMessage(ctx.chatHistory, 'system', 'Auto-approval enabled for tasks'), + }), + }, + [AIChatMachineEventType.DISABLE_AUTO_APPROVE]: { + actions: assign({ + autoApproveEnabled: (_ctx) => false, + chatHistory: (ctx) => + addChatMessage(ctx.chatHistory, 'system', 'Auto-approval disabled for tasks'), + }), + }, [AIChatMachineEventType.RESET]: { target: 'Idle', actions: [ diff --git a/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts b/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts index 7b344b5eaa9..9d6df0e92cf 100644 --- a/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts +++ b/workspaces/ballerina/ballerina-rpc-client/src/BallerinaRpcClient.ts @@ -56,7 +56,9 @@ import { AIChatMachineEventType, aiChatStateChanged, AIChatMachineSendableEvent, - AIChatMachineStateValue + AIChatMachineStateValue, + getAIChatContext, + AIChatMachineContext } from "@wso2/ballerina-core"; import { LangClientRpcClient } from "./rpc-clients/lang-client/rpc-client"; import { LibraryBrowserRpcClient } from "./rpc-clients/library-browser/rpc-client"; @@ -208,6 +210,10 @@ export class BallerinaRpcClient { this.messenger.sendRequest(sendAIChatStateEvent, HOST_EXTENSION, event); } + getAIChatContext(): Promise { + return this.messenger.sendRequest(getAIChatContext, HOST_EXTENSION); + } + onProjectContentUpdated(callback: (state: boolean) => void) { this.messenger.onNotification(projectContentUpdated, callback); } 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 09770049797..b128c15479c 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 @@ -153,6 +153,7 @@ const AIChat: React.FC = () => { const [showSettings, setShowSettings] = useState(false); const [aiChatStateMachineState, setAiChatStateMachineState] = useState("Idle"); + const [isAutoApproveEnabled, setIsAutoApproveEnabled] = useState(false); // Approval dialog state const [approvalRequest, setApprovalRequest] = useState<{ @@ -290,6 +291,23 @@ const AIChat: React.FC = () => { }); }, []); + // Initialize auto-approve state from state machine context on mount + useEffect(() => { + const initializeAutoApproveState = async () => { + try { + const context = await rpcClient.getAIChatContext(); + if (context && context.autoApproveEnabled !== undefined) { + setIsAutoApproveEnabled(context.autoApproveEnabled); + console.log(`[AIChat] Initialized auto-approve state: ${context.autoApproveEnabled}`); + } + } catch (error) { + console.error("[AIChat] Failed to initialize auto-approve state:", error); + } + }; + + initializeAutoApproveState(); + }, [rpcClient]); + rpcClient?.onAIChatStateChanged((newState: AIChatMachineStateValue) => { setAiChatStateMachineState(newState); }); @@ -1885,6 +1903,19 @@ const AIChat: React.FC = () => { rpcClient.sendAIChatStateEvent(AIChatMachineEventType.RESET); } + // Handle auto-approval toggle + const handleToggleAutoApprove = () => { + const newValue = !isAutoApproveEnabled; + setIsAutoApproveEnabled(newValue); + + // Send event to state machine + if (newValue) { + rpcClient.sendAIChatStateEvent(AIChatMachineEventType.ENABLE_AUTO_APPROVE); + } else { + rpcClient.sendAIChatStateEvent(AIChatMachineEventType.DISABLE_AUTO_APPROVE); + } + }; + const questionMessages = messages.filter((message) => message.type === "question"); if (questionMessages.length > 0) { localStorage.setItem( @@ -2155,6 +2186,14 @@ const AIChat: React.FC = () => {
State: {aiChatStateMachineState}
+ - {/* Show TodoSection when we have tasks from approval request OR after plan is approved */} {(approvalRequest?.tasks || todoTasks) && ( = ({ tasks, message, onApprove, on const hasInProgress = !!inProgressTask; const needsApproval = !!approvalType; - // Highlight container only for plan approval, not task completion approval useEffect(() => { - if (approvalType === "plan") { - // Highlight entire container for plan approval - setIsNew(true); - } else { - // No container highlight for task approval (task itself has blue border) - setIsNew(false); - } - - // Reset submission state when approval type changes + setIsNew(approvalType === "plan"); setIsSubmitting(false); setShowRejectComment(false); setRejectComment(""); @@ -280,7 +271,6 @@ const TodoSection: React.FC = ({ tasks, message, onApprove, on setIsExpanded(!isExpanded); }; - // Function to scroll to in-progress task const scrollToInProgress = () => { if (inProgressRef.current) { inProgressRef.current.scrollIntoView({ @@ -290,25 +280,21 @@ const TodoSection: React.FC = ({ tasks, message, onApprove, on } }; - // Auto-scroll to in-progress task useEffect(() => { if (isExpanded && hasInProgress) { scrollToInProgress(); } }, [isExpanded, inProgressTask?.description]); - // Handle user scroll - refocus after delay useEffect(() => { const todoList = todoListRef.current; if (!todoList || !hasInProgress) return; const handleScroll = () => { - // Clear existing timeout if (scrollTimeoutRef.current) { clearTimeout(scrollTimeoutRef.current); } - // Set new timeout to refocus after 3 seconds scrollTimeoutRef.current = setTimeout(() => { scrollToInProgress(); }, 3000); @@ -324,7 +310,6 @@ const TodoSection: React.FC = ({ tasks, message, onApprove, on }; }, [hasInProgress, inProgressTask?.description]); - // Determine status text const getStatusText = () => { if (allCompleted) return "completed"; if (needsApproval) return "awaiting approval"; From 2f48ea55341f74e1e7f38a8469be630651cdaaf8 Mon Sep 17 00:00:00 2001 From: RNViththagan Date: Thu, 30 Oct 2025 12:04:33 +0530 Subject: [PATCH 015/619] Refactor approval UI into separate footer component Extracts approval functionality into dedicated component for better organization and reusability. - Create new ApprovalFooter component for handling task and plan approvals - Extract approval logic from TodoSection to dedicated footer component - Refactor AI chat machine to simplify state management - Update AI chat component to use new approval footer structure - Improve separation of concerns between todo display and approval actions - Reduce TodoSection complexity by removing embedded approval UI - Enhance code maintainability with modular approval component structure --- .../src/views/ai-panel/aiChatMachine.ts | 46 +--- .../AIChat/Footer/ApprovalFooter.tsx | 244 ++++++++++++++++++ .../views/AIPanel/components/AIChat/index.tsx | 64 +++-- .../views/AIPanel/components/TodoSection.tsx | 191 +------------- 4 files changed, 298 insertions(+), 247 deletions(-) create mode 100644 workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/ApprovalFooter.tsx diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts index f2e3234bd1b..145403e3c1f 100644 --- a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts +++ b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts @@ -86,6 +86,17 @@ const chatMachine = createMachine event.payload.prompt, + chatHistory: (ctx, event) => + addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), + currentPlan: (_ctx) => undefined, + currentTaskIndex: (_ctx) => -1, + errorMessage: (_ctx) => undefined, + }), + }, [AIChatMachineEventType.ENABLE_AUTO_APPROVE]: { actions: assign({ autoApproveEnabled: (_ctx) => true, @@ -141,17 +152,6 @@ const chatMachine = createMachine generateSessionId(), projectId: (_ctx) => generateProjectId(), }), - on: { - [AIChatMachineEventType.SUBMIT_PROMPT]: { - target: 'Initiating', - actions: assign({ - initialPrompt: (_ctx, event) => event.payload.prompt, - chatHistory: (ctx, event) => - addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), - errorMessage: (_ctx) => undefined, - }), - }, - }, }, //TODO : Optional state we can remove if not needed. its just to show that generation is starting. Initiating: { @@ -436,34 +436,10 @@ const chatMachine = createMachine event.payload.prompt, - chatHistory: (ctx, event) => - addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), - currentPlan: (_ctx) => undefined, - currentTaskIndex: (_ctx) => -1, - errorMessage: (_ctx) => undefined, - }), - }, - }, }, PartiallyCompleted: { entry: 'saveChatState', on: { - [AIChatMachineEventType.SUBMIT_PROMPT]: { - target: 'Initiating', - actions: assign({ - initialPrompt: (_ctx, event) => event.payload.prompt, - chatHistory: (ctx, event) => - addChatMessage(ctx.chatHistory, 'user', event.payload.prompt), - currentPlan: (_ctx) => undefined, - currentTaskIndex: (_ctx) => -1, - errorMessage: (_ctx) => undefined, - }), - }, [AIChatMachineEventType.START_TASK_EXECUTION]: { target: 'ExecutingTask', actions: assign({ diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/ApprovalFooter.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/ApprovalFooter.tsx new file mode 100644 index 00000000000..bd06cd7968c --- /dev/null +++ b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/ApprovalFooter.tsx @@ -0,0 +1,244 @@ +/** + * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com) All Rights Reserved. + * + * WSO2 LLC. licenses this file to you under the Apache License, + * Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import React, { useState, useEffect } from "react"; +import styled from "@emotion/styled"; +import { FooterContainer } from "./index"; + +const ApprovalContainer = styled.div` + display: flex; + flex-direction: column; + gap: 12px; +`; + +const PromptText = styled.div` + font-size: 13px; + font-weight: 500; + color: var(--vscode-editor-foreground); + margin-bottom: 4px; +`; + +const ButtonsColumn = styled.div` + display: flex; + flex-direction: column; + gap: 8px; +`; + +const Button = styled.button<{ variant?: "primary" | "secondary" }>` + padding: 8px 14px; + font-size: 12px; + font-weight: 500; + border: none; + border-radius: 4px; + cursor: pointer; + transition: all 0.2s ease; + font-family: var(--vscode-editor-font-family); + white-space: nowrap; + width: 100%; + text-align: left; + + ${(props: { variant?: "primary" | "secondary" }) => + props.variant === "primary" + ? ` + background-color: var(--vscode-button-background); + color: var(--vscode-button-foreground); + + &:hover { + background-color: var(--vscode-button-hoverBackground); + } + ` + : ` + background-color: var(--vscode-button-secondaryBackground); + color: var(--vscode-button-secondaryForeground); + + &:hover { + background-color: var(--vscode-button-secondaryHoverBackground); + } + `} + + &:active { + opacity: 0.8; + } + + &:disabled { + opacity: 0.5; + cursor: not-allowed; + } +`; + +const InputContainer = styled.div` + display: flex; + align-items: center; + gap: 6px; + width: 100%; + background-color: var(--vscode-input-background); + border: 1px solid var(--vscode-input-border); + border-radius: 4px; + padding: 8px 14px; + transition: border-color 0.2s ease; + height: 36px; + box-sizing: border-box; + + &:focus-within { + border-color: var(--vscode-focusBorder); + } +`; + +const Input = styled.input` + flex: 1; + background: transparent; + border: none; + color: var(--vscode-input-foreground); + font-family: var(--vscode-editor-font-family); + font-size: 12px; + outline: none; + padding: 0; + + &::placeholder { + color: var(--vscode-input-placeholderForeground); + } +`; + +const SendButton = styled.button` + background: transparent; + border: none; + color: var(--vscode-icon-foreground); + cursor: pointer; + padding: 4px; + display: flex; + align-items: center; + justify-content: center; + border-radius: 3px; + transition: background-color 0.2s ease; + + &:hover:not(:disabled) { + background-color: var(--vscode-toolbar-hoverBackground); + } + + &:disabled { + opacity: 0.4; + cursor: not-allowed; + } +`; + +type ApprovalType = "plan" | "completion"; + +interface ApprovalFooterProps { + approvalType: ApprovalType; + onApprove: (enableAutoApprove: boolean) => void; + onReject: (comment: string) => void; + isSubmitting?: boolean; +} + +const ApprovalFooter: React.FC = ({ + approvalType, + onApprove, + onReject, + isSubmitting = false, +}) => { + const [comment, setComment] = useState(""); + + useEffect(() => { + // Reset comment when approval type changes + setComment(""); + }, [approvalType]); + + const handleApproveWithAutoApprove = () => { + onApprove(true); + }; + + const handleApproveManually = () => { + onApprove(false); + }; + + const handleRejectSubmit = () => { + const trimmedComment = comment.trim(); + if (trimmedComment) { + onReject(trimmedComment); + setComment(""); + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter") { + e.preventDefault(); + if (comment.trim()) { + handleRejectSubmit(); + } + } + }; + + const handleInputChange = (e: React.ChangeEvent) => { + setComment(e.target.value); + }; + + const promptText = approvalType === "plan" + ? "Accept this plan?" + : "Approve this task?"; + + const primaryButtonText = approvalType === "plan" + ? "Yes and auto-approve tasks" + : "Yes and don't ask again"; + + const secondaryButtonText = "Yes"; + + const placeholderText = "Or describe what needs to change..."; + + return ( + + + {promptText} + + + + + + + + + + + + + ); +}; + +export default ApprovalFooter; 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 d65cd8b0e1a..fb1b7bf9b81 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 @@ -85,6 +85,7 @@ import { SYSTEM_ERROR_SECRET } from "../AIChatInput/constants"; import { CodeSegment } from "../CodeSegment"; import AttachmentBox, { AttachmentsContainer } from "../AttachmentBox"; import Footer from "./Footer"; +import ApprovalFooter from "./Footer/ApprovalFooter"; import { useFooterLogic } from "./Footer/useFooterLogic"; import { SettingsPanel } from "../../SettingsPanel"; import WelcomeMessage from "./Welcome"; @@ -2077,13 +2078,21 @@ const AIChat: React.FC = () => { }); }; - const handleApprovalApprove = (comment?: string) => { + const handleApprovalApprove = (enableAutoApprove: boolean) => { if (!approvalRequest) return; + if (enableAutoApprove && !isAutoApproveEnabled) { + setIsAutoApproveEnabled(true); + rpcClient.sendAIChatStateEvent(AIChatMachineEventType.ENABLE_AUTO_APPROVE); + } else if (!enableAutoApprove && isAutoApproveEnabled) { + setIsAutoApproveEnabled(false); + rpcClient.sendAIChatStateEvent(AIChatMachineEventType.DISABLE_AUTO_APPROVE); + } + if (approvalRequest.approvalType === "plan") { rpcClient.sendAIChatStateEvent({ type: AIChatMachineEventType.APPROVE_PLAN, - payload: { comment } + payload: {} }); setTodoTasks(approvalRequest.tasks); setTasksMessage(approvalRequest.message); @@ -2095,7 +2104,6 @@ const AIChat: React.FC = () => { rpcClient.sendAIChatStateEvent({ type: AIChatMachineEventType.APPROVE_TASK, payload: { - comment, lastApprovedTaskIndex: lastApprovedTaskIndex >= 0 ? lastApprovedTaskIndex : undefined } }); @@ -2104,7 +2112,7 @@ const AIChat: React.FC = () => { setApprovalRequest(null); }; - const handleApprovalReject = (comment?: string) => { + const handleApprovalReject = (comment: string) => { if (!approvalRequest) return; if (approvalRequest.approvalType === "plan") { @@ -2163,9 +2171,6 @@ const AIChat: React.FC = () => { )} @@ -2414,24 +2419,33 @@ const AIChat: React.FC = () => { })}
-