2020import { createMachine , assign , interpret } from 'xstate' ;
2121import { extension } from '../../BalExtensionContext' ;
2222import * as crypto from 'crypto' ;
23- import { AIChatMachineContext , AIChatMachineEventType , AIChatMachineSendableEvent , AIChatMachineStateValue , ChatMessage , Plan , Task , TaskStatus , UIChatHistoryMessage } from '@wso2/ballerina-core/lib/state-machine-types' ;
23+ import { AIChatMachineContext , AIChatMachineEventType , AIChatMachineSendableEvent , AIChatMachineStateValue , ChatMessage , Plan , Task , TaskStatus , UIChatHistoryMessage , Checkpoint } from '@wso2/ballerina-core/lib/state-machine-types' ;
2424import { workspace } from 'vscode' ;
2525import { GenerateAgentCodeRequest } from '@wso2/ballerina-core/lib/rpc-types/ai-panel/interfaces' ;
2626import { generateDesign } from '../../features/ai/service/design/design' ;
27+ import { captureWorkspaceSnapshot , restoreWorkspaceSnapshot } from './checkpoint/checkpointUtils' ;
28+ import { getCheckpointConfig } from './checkpoint/checkpointConfig' ;
29+ import { notifyCheckpointCaptured } from '../../RPCLayer' ;
2730
2831const generateId = ( ) => `${ Date . now ( ) } ${ Math . random ( ) . toString ( 36 ) . substr ( 2 , 9 ) } ;
2932
@@ -93,6 +96,67 @@ const updateChatMessage = (
9396 } ) ;
9497} ;
9598
99+ const cleanupOldCheckpoints = ( checkpoints : Checkpoint [ ] ) : Checkpoint [ ] => {
100+ const config = getCheckpointConfig ( ) ;
101+ if ( checkpoints . length <= config . maxCount ) {
102+ return checkpoints ;
103+ }
104+ return checkpoints . slice ( - config . maxCount ) ;
105+ } ;
106+
107+ const captureCheckpointAction = ( context : AIChatMachineContext ) => {
108+ const lastMessage = context . chatHistory [ context . chatHistory . length - 1 ] ;
109+ if ( ! lastMessage ) {
110+ return ;
111+ }
112+
113+ captureWorkspaceSnapshot ( lastMessage . id ) . then ( checkpoint => {
114+ if ( checkpoint ) {
115+ lastMessage . checkpointId = checkpoint . id ;
116+ const updatedCheckpoints = cleanupOldCheckpoints ( [ ...( context . checkpoints || [ ] ) , checkpoint ] ) ;
117+ context . checkpoints = updatedCheckpoints ;
118+ saveChatState ( context ) ;
119+
120+ // Notify frontend that checkpoint is captured
121+ notifyCheckpointCaptured ( {
122+ messageId : lastMessage . id ,
123+ checkpointId : checkpoint . id
124+ } ) ;
125+ }
126+ } ) . catch ( error => {
127+ console . error ( '[Checkpoint] Failed to capture checkpoint:' , error ) ;
128+ } ) ;
129+ } ;
130+
131+ const restoreCheckpointAction = ( context : AIChatMachineContext , event : any ) => {
132+ const checkpointId = event . payload . checkpointId ;
133+ const checkpoint = context . checkpoints ?. find ( c => c . id === checkpointId ) ;
134+
135+ if ( ! checkpoint ) {
136+ console . error ( `[Checkpoint] Checkpoint ${ checkpointId } ) ;
137+ return ;
138+ }
139+
140+ const messageIndex = context . chatHistory . findIndex ( m => m . id === checkpoint . messageId ) ;
141+ const restoredHistory = messageIndex >= 0 ? context . chatHistory . slice ( 0 , messageIndex ) : context . chatHistory ;
142+
143+ const checkpointIndex = context . checkpoints ?. findIndex ( c => c . id === checkpointId ) || 0 ;
144+ const restoredCheckpoints = checkpointIndex >= 0 ? ( context . checkpoints ?. slice ( 0 , checkpointIndex ) || [ ] ) : ( context . checkpoints || [ ] ) ;
145+
146+ context . chatHistory = restoredHistory ;
147+ context . checkpoints = restoredCheckpoints ;
148+ context . currentPlan = undefined ;
149+ context . currentTaskIndex = - 1 ;
150+
151+ saveChatState ( context ) ;
152+
153+ restoreWorkspaceSnapshot ( checkpoint ) . then ( ( ) => {
154+
155+ } ) . catch ( error => {
156+ console . error ( '[Checkpoint] Failed to restore workspace snapshot:' , error ) ;
157+ } ) ;
158+ } ;
159+
96160const chatMachine = createMachine < AIChatMachineContext , AIChatMachineSendableEvent > ( {
97161 /** @xstate -layout N4IgpgJg5mDOIC5QCMCGAbdYBOBLAdqgLSq5EDGAFqgC4B0AwtmLQVAArqr4DEEA9vjB0CAN34BrYeWa0wnbgG0ADAF1EoAA79YuGrkEaQAD0QBmACxm6ZgJzKAbACZbTpwEYnZ97YA0IAE9EAHZg9zoADk8HZQBWW2DHC1iLAF9U-zRMHAJiUgpqeiYWfXwOLl4cbH5sOk0uGgAzGoBbOhkS+QqVdSQQbV19Qz7TBEtrO0cXN09vP0DECIdYuidl4KcLdwdgiPtY9MyMLDxCEjIqWjoAVU0IVjKFXgEhEXxxKToAVzu5J56jAM9AZ8EZRuMbPZnK4PF4fP4gghYhELHQoWYHLZYsEHO5YrEzIcQFkTrlzgUrrd7qVytweFUanUGs1sG0ftSuko1ICdMDhqBwVZIVMYbN4QsEDsIjZ4hYIqFgpZlMo0hlicccmd8pd6ABRYxgchfGkAFVQsAkfEEwjEkmEYANRpoYDNFoBfSBQ1BI3MQsm0JmcPmiKidAstgjtnD9lcZkcRJJmryF0KdH1huNbFdloZtXqtBZbQdGed2fdWl5XrBvomUOmsLmCMQ7mCqKx3nc3giUQcUWCCY1p2TFPoADECGAaF98FnzZaXjb3na6I0J1OhGXuR7KyDq2M-XXRUGm2NcZF4pYsRE1sjVUdskPyTq6OOhOvZxb6dhqnnma0V2u04unO5b9Du-ImDWwoBg24qIms0oWFM3hmBEKSdoSRL4PwEBwEYiaPtqhQ8oMu4+ggRBmNYiSxEkeLuHsSGxCeRArJGEZOPKtF2O4cQDg+ZJEVcACSEBYCRfLegKzadqsDHBLE7jbE42JKSe4ZOHQnacS416Yg4yQHGqBGCSmVzFA8tJSWBpEQaM7ioi44ZmLeURxFiJ44ppyqxF4yhLDETitvxpJamZ9BPAASmAoi4GAADuElVuR6E2FYOIRGYwRQnBiDxIhDjOIqqH2N2d7qgJYUjjcvw0k8SVkdJSK2GiUwWBYQVbBsYQnvEwR0A4ZhOT52L4u4IVJk+qbpk6H4SA1dmIMkdAEvszgMVGjhmCeaErfl14ooNPhWBNhHhXQACC8WkDS1ywDgF3kMlFa2dZoypeMGVZTlwbNteYbZb2sp2IN+KnaZ1WvpO05zQtb1LU4J6ceE2kKVlWy8di4NVc+ACKXxwF6bB3TgcN7iiNhDcoazZT4g2OCeqEOANQ1LCDth7INmH3qFw7Pgw-AtPUk6QGT5EU1RTjUzsth03GDg7fYdDKDxHVxoqnZGTzk1CXq341GLTUfelvbfY4uUIIk-XIjEVEdTe8bpKkQA */
98162 id : 'ballerina-ai-chat' ,
@@ -110,17 +174,21 @@ const chatMachine = createMachine<AIChatMachineContext, AIChatMachineSendableEve
110174 autoApproveEnabled : false ,
111175 previousState : undefined ,
112176 currentSpec : undefined ,
177+ checkpoints : [ ] ,
113178 } as AIChatMachineContext ,
114179 on : {
115180 [ AIChatMachineEventType . SUBMIT_PROMPT ] : {
116181 target : 'Initiating' ,
117- actions : assign ( {
118- chatHistory : ( ctx , event ) =>
119- addUserMessage ( ctx . chatHistory , event . payload . prompt ) ,
120- currentPlan : ( _ctx ) => undefined ,
121- currentTaskIndex : ( _ctx ) => - 1 ,
122- errorMessage : ( _ctx ) => undefined ,
123- } ) ,
182+ actions : [
183+ assign ( {
184+ chatHistory : ( ctx , event ) =>
185+ addUserMessage ( ctx . chatHistory , event . payload . prompt ) ,
186+ currentPlan : ( _ctx ) => undefined ,
187+ currentTaskIndex : ( _ctx ) => - 1 ,
188+ errorMessage : ( _ctx ) => undefined ,
189+ } ) ,
190+ 'captureCheckpoint' ,
191+ ] ,
124192 } ,
125193 [ AIChatMachineEventType . UPDATE_CHAT_MESSAGE ] : {
126194 actions : assign ( {
@@ -151,7 +219,7 @@ const chatMachine = createMachine<AIChatMachineContext, AIChatMachineSendableEve
151219 currentQuestion : ( _ctx ) => undefined ,
152220 errorMessage : ( _ctx ) => undefined ,
153221 sessionId : ( _ctx ) => undefined ,
154- // Keep projectId to maintain project context
222+ checkpoints : ( _ctx ) => [ ] ,
155223 } ) ,
156224 ] ,
157225 } ,
@@ -166,6 +234,10 @@ const chatMachine = createMachine<AIChatMachineContext, AIChatMachineSendableEve
166234 sessionId : ( _ctx , event ) => event . payload . state . sessionId ,
167235 } ) ,
168236 } ,
237+ [ AIChatMachineEventType . RESTORE_CHECKPOINT ] : {
238+ target : 'Idle' ,
239+ actions : [ 'restoreCheckpoint' ] ,
240+ } ,
169241 [ AIChatMachineEventType . ERROR ] : {
170242 target : 'Error' ,
171243 actions : assign ( {
@@ -547,7 +619,9 @@ const convertChatHistoryToUIMessages = (chatHistory: ChatMessage[]): UIChatHisto
547619
548620 messages . push ( {
549621 role : 'user' ,
550- content : msg . content
622+ content : msg . content ,
623+ checkpointId : msg . checkpointId ,
624+ messageId : msg . id
551625 } ) ;
552626
553627 if ( msg . uiResponse ) {
@@ -614,6 +688,7 @@ const saveChatState = (context: AIChatMachineContext) => {
614688 currentTaskIndex : context . currentTaskIndex ,
615689 sessionId : context . sessionId ,
616690 projectId : context . projectId ,
691+ checkpoints : context . checkpoints || [ ] ,
617692 savedAt : Date . now ( ) ,
618693 } ;
619694
@@ -764,6 +839,8 @@ const chatStateService = interpret(
764839 actions : {
765840 saveChatState : ( context ) => saveChatState ( context ) ,
766841 clearChatState : ( context ) => clearChatStateAction ( context ) ,
842+ captureCheckpoint : ( context ) => captureCheckpointAction ( context ) ,
843+ restoreCheckpoint : ( context , event ) => restoreCheckpointAction ( context , event ) ,
767844 } ,
768845 } )
769846) ;
0 commit comments