Merge pull request #992 from RNViththagan/copilot-agent-checkpoint-mechanism

Implement checkpoint mechanism for AI copilot chat

Author: xlight05

common/config/rush/pnpm-lock.yaml

@@ -421,6 +421,7 @@ export enum AIChatMachineEventType {
     DISABLE_AUTO_APPROVE = 'DISABLE_AUTO_APPROVE',
     REJECT_TASK = 'REJECT_TASK',
     RESTORE_STATE = 'RESTORE_STATE',
+    RESTORE_CHECKPOINT = 'RESTORE_CHECKPOINT',
     ERROR = 'ERROR',
     RETRY = 'RETRY',
     CONNECTOR_GENERATION_REQUESTED = 'CONNECTOR_GENERATION_REQUESTED',
@@ -434,6 +435,16 @@ export interface ChatMessage {
     uiResponse: string;
     modelMessages: any[];
     timestamp: number;
+    checkpointId?: string;
+}
+
+export interface Checkpoint {
+    id: string;
+    messageId: string;
+    timestamp: number;
+    workspaceSnapshot: { [filePath: string]: string };
+    fileList: string[];
+    snapshotSize: number;
 }
 
 /**
@@ -497,6 +508,7 @@ export interface AIChatMachineContext {
         comment?: string;
     };
     previousState?: AIChatMachineStateValue;
+    checkpoints?: Checkpoint[];
 }
 
 export type AIChatMachineSendableEvent =
@@ -515,6 +527,7 @@ export type AIChatMachineSendableEvent =
     | { type: AIChatMachineEventType.REJECT_TASK; payload: { comment?: string } }
     | { type: AIChatMachineEventType.RESET }
     | { type: AIChatMachineEventType.RESTORE_STATE; payload: { state: AIChatMachineContext } }
+    | { type: AIChatMachineEventType.RESTORE_CHECKPOINT; payload: { checkpointId: string } }
     | { type: AIChatMachineEventType.ERROR; payload: { message: string } }
     | { type: AIChatMachineEventType.RETRY }
     | { type: AIChatMachineEventType.CONNECTOR_GENERATION_REQUESTED; payload: { requestId: string; serviceName?: string; serviceDescription?: string; fromState?: AIChatMachineStateValue } }
@@ -577,6 +590,8 @@ export enum ColorThemeKind {
 export interface UIChatHistoryMessage {
     role: "user" | "assistant";
     content: string;
+    checkpointId?: string;
+    messageId?: string;
 }
 
 export const aiStateChanged: NotificationType<AIMachineStateValue> = { method: 'aiStateChanged' };
@@ -588,6 +603,12 @@ export const sendAIChatStateEvent: RequestType<AIChatMachineEventType | AIChatMa
 export const getAIChatContext: RequestType<void, AIChatMachineContext> = { method: 'getAIChatContext' };
 export const getAIChatUIHistory: RequestType<void, UIChatHistoryMessage[]> = { method: 'getAIChatUIHistory' };
 
+export interface CheckpointCapturedPayload {
+    messageId: string;
+    checkpointId: string;
+}
+export const checkpointCaptured: NotificationType<CheckpointCapturedPayload> = { method: 'checkpointCaptured' };
+
 // Connector Generator RPC methods
 export interface ConnectorGeneratorResponsePayload {
     requestId: string;

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

@@ -233,6 +233,33 @@
                         "Trace request/response messages only",
                         "Trace request/response messages with parameters and content"
                     ]
+                },
+                "ballerina.copilot.checkpoints.enabled": {
+                    "type": "boolean",
+                    "default": true,
+                    "description": "Enable checkpoint mechanism for copilot chat to restore previous states."
+                },
+                "ballerina.copilot.checkpoints.maxCount": {
+                    "type": "number",
+                    "default": 20,
+                    "description": "Maximum number of checkpoints to keep."
+                },
+                "ballerina.copilot.checkpoints.ignorePatterns": {
+                    "type": "array",
+                    "default": [
+                        ".git/**",
+                        "target/**",
+                        "build/**",
+                        "dist/**",
+                        ".vscode/**",
+                        "*.log"
+                    ],
+                    "description": "File patterns to ignore when creating checkpoints."
+                },
+                "ballerina.copilot.checkpoints.maxSnapshotSize": {
+                    "type": "number",
+                    "default": 52428800,
+                    "description": "Maximum snapshot size in bytes (default: 50MB)."
                 }
             }
         },
@@ -1198,8 +1225,8 @@
         "copyJSLibs": "copyfiles -f ../ballerina-visualizer/build/*.js resources/jslibs && copyfiles -f ../trace-visualizer/build/*.js resources/jslibs"
     },
     "dependencies": {
-        "@ai-sdk/amazon-bedrock": "^3.0.25",
-        "@ai-sdk/anthropic": "^2.0.20",
+        "@ai-sdk/amazon-bedrock": "^3.0.57",
+        "@ai-sdk/anthropic": "^2.0.45",
         "@opentelemetry/auto-instrumentations-node": "^0.67.0",
         "@opentelemetry/sdk-node": "^0.208.0",
         "@types/lodash": "^4.14.200",
@@ -1211,7 +1238,7 @@
         "@wso2/syntax-tree": "workspace:*",
         "@wso2/trace-visualizer": "workspace:*",
         "@wso2/wso2-platform-core": "workspace:*",
-        "ai": "^5.0.56",
+        "ai": "^5.0.101",
         "cors-anywhere": "^0.4.4",
         "del-cli": "^5.1.0",
         "dotenv": "~16.5.0",

workspaces/ballerina/ballerina-extension/package.json

@@ -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, getAIChatContext, getAIChatUIHistory, 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, getAIChatUIHistory, AIChatMachineEventType, AIChatMachineSendableEvent, checkpointCaptured, CheckpointCapturedPayload } 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';
@@ -190,3 +190,7 @@ export function notifyAiWebview() {
 export function notifyBreakpointChange() {
     RPCLayer._messenger.sendNotification(breakpointChanged, { type: 'webview', webviewType: VisualizerWebview.viewType }, true);
 }
+
+export function notifyCheckpointCaptured(payload: CheckpointCapturedPayload) {
+    RPCLayer._messenger.sendNotification(checkpointCaptured, { type: 'webview', webviewType: AiPanelWebview.viewType }, payload);
+}