[AI Datamapper] Fix check error handling and output formatting

workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts

@@ -121,7 +121,7 @@ export interface ProjectImports {
 // Data-mapper related interfaces
 export interface MetadataWithAttachments {
     metadata: ExtendedDataMapperMetadata;
-    attachments?: Attachment[];
+    attachments: Attachment[];
 }
 
 export interface InlineMappingsSourceResult {

workspaces/ballerina/ballerina-extension/src/features/ai/dataMapping.ts

@@ -33,6 +33,7 @@ import { PackageInfo, TypesGenerationResult } from "./service/datamapper/types";
 import { URI } from "vscode-uri";
 import { getAllDataMapperSource } from "./service/datamapper/datamapper";
 import { StateMachine } from "../../stateMachine";
+import { CopilotEventHandler } from "./service/event";
 
 // Set to false to include mappings with default values
 const OMIT_DEFAULT_MAPPINGS_ENABLED = true;
@@ -473,13 +474,15 @@ export async function createTempFileAndGenerateMetadata(params: CreateTempFileRe
 
 export async function generateMappings(
   metadataWithAttachments: MetadataWithAttachments,
-  context: any
+  context: any,
+  eventHandler: CopilotEventHandler
 ): Promise<AllDataMapperSourceRequest> {
   const targetFilePath = metadataWithAttachments.metadata.codeData.lineRange.fileName || context.documentUri;
 
   const generatedMappings = await generateMappingExpressionsFromModel(
     metadataWithAttachments.metadata.mappingsModel as DMModel,
-    metadataWithAttachments.attachments || []
+    metadataWithAttachments.attachments || [],
+    eventHandler
   );
 
   const customFunctionMappings = generatedMappings.filter(mapping => mapping.isFunctionCall);
@@ -1060,6 +1063,7 @@ export async function generateInlineMappingsSource(
   inlineMappingRequest: MetadataWithAttachments,
   langClient: ExtendedLangClient,
   context: any,
+  eventHandler: CopilotEventHandler
 ): Promise<InlineMappingsSourceResult> {
   if (!inlineMappingRequest) {
     throw new Error("Inline mapping request is required");
@@ -1096,16 +1100,18 @@ export async function generateInlineMappingsSource(
 
   // Prepare mapping request payload
   const mappingRequestPayload: MetadataWithAttachments = {
-    metadata: tempFileMetadata
+    metadata: tempFileMetadata,
+    attachments: []
   };
-  if (inlineMappingRequest.attachments && inlineMappingRequest.attachments.length > 0) {
+  if (inlineMappingRequest.attachments.length > 0) {
     mappingRequestPayload.attachments = inlineMappingRequest.attachments;
   }
 
   // Generate mappings and source code
   const allMappingsRequest = await generateMappings(
     mappingRequestPayload,
-    context
+    context,
+    eventHandler
   );
 
   const generatedSourceResponse = await getAllDataMapperSource(allMappingsRequest);

workspaces/ballerina/ballerina-extension/src/features/ai/service/datamapper/types.ts

@@ -53,6 +53,11 @@ export interface RepairedFiles {
     repairedFiles: SourceFile[];
 }
 
+export interface CodeRepairResult {
+    finalContent: string;
+    customFunctionsContent: string;
+}
+
 // =============================================================================
 // MAPPING HINTS
 // =============================================================================

workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/repair-utils.ts

@@ -330,3 +330,98 @@ export async function addMissingRequiredFields(
 
     return projectModified;
 }
+
+export async function addCheckExpressionErrors(
+    diagnosticsResult: Diagnostics[],
+    langClient: ExtendedLangClient
+): Promise<boolean> {
+    let projectModified = false;
+
+    for (const diag of diagnosticsResult) {
+        const fileUri = diag.uri;
+        const diagnostics = diag.diagnostics;
+
+        // Filter BCE3032 diagnostics (check expression errors)
+        const checkExprDiagnostics = diagnostics.filter(d => d.code === "BCE3032");
+        if (!checkExprDiagnostics.length) {
+            continue;
+        }
+
+        const astModifications: STModification[] = [];
+
+        // Process each diagnostic individually
+        for (const diagnostic of checkExprDiagnostics) {
+            try {
+                // Get code actions for the diagnostic
+                const codeActions = await langClient.codeAction({
+                    textDocument: { uri: fileUri },
+                    range: {
+                        start: diagnostic.range.start,
+                        end: diagnostic.range.end
+                    },
+                    context: {
+                        diagnostics: [diagnostic],
+                        only: ['quickfix'],
+                        triggerKind: 1
+                    }
+                });
+
+                if (!codeActions?.length) {
+                    console.warn(`No code actions returned for ${fileUri} at line ${diagnostic.range.start.line}`);
+                    continue;
+                }
+
+                // Find the action that adds error to return type
+                // The language server typically provides actions like "Change return type to ..."
+                const action = codeActions.find(
+                    action => action.title && (
+                        action.title.toLowerCase().includes("change") &&
+                        action.title.toLowerCase().includes("return") &&
+                        action.title.toLowerCase().includes("error")
+                    )
+                );
+
+                if (!action?.edit?.documentChanges?.length) {
+                    continue;
+                }
+
+                const docEdit = action.edit.documentChanges[0] as TextDocumentEdit;
+
+                // Process all edits from the code action
+                for (const edit of docEdit.edits) {
+                    astModifications.push({
+                        startLine: edit.range.start.line,
+                        startColumn: edit.range.start.character,
+                        endLine: edit.range.end.line,
+                        endColumn: edit.range.end.character,
+                        type: "INSERT",
+                        isImport: false,
+                        config: { STATEMENT: edit.newText }
+                    });
+                }
+            } catch (err) {
+                console.warn(`Could not apply code action for ${fileUri} at line ${diagnostic.range.start.line}:`, err);
+            }
+        }
+
+        // Apply modifications to syntax tree
+        if (astModifications.length > 0) {
+            const syntaxTree = await langClient.stModify({
+                documentIdentifier: { uri: fileUri },
+                astModifications: astModifications
+            });
+
+            // Update file content
+            const { source } = syntaxTree as SyntaxTree;
+            if (!source) {
+                // Handle the case where source is undefined, when compiler issue occurs
+                return false;
+            }
+            const absolutePath = fileURLToPath(fileUri);
+            writeBallerinaFileDidOpenTemp(absolutePath, source);
+            projectModified = true;
+        }
+    }
+
+    return projectModified;
+}

workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts

@@ -673,7 +673,23 @@ export class AiPanelRpcManager implements AIPanelAPI {
                 throw new Error("Not a Ballerina project.");
             }
             await addToIntegration(projectPath, params.fileChanges);
-            updateView();
+
+            const context = StateMachine.context();
+            const dataMapperMetadata = context.dataMapperMetadata;
+            if (!dataMapperMetadata || !dataMapperMetadata.codeData) {
+                updateView();
+                return true;
+            }
+
+            // Refresh data mapper with the updated code
+            let filePath = dataMapperMetadata.codeData.lineRange?.fileName;
+            const varName = dataMapperMetadata.name;
+            if (!filePath || !varName) {
+                updateView();
+                return true;
+            }
+
+            await refreshDataMapper(filePath, dataMapperMetadata.codeData, varName);
             return true;
         } catch (error) {
             console.error(">>> Failed to add files to the project", error);

workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/utils.ts

@@ -28,6 +28,7 @@ import { getAskResponse } from "../../../src/features/ai/service/ask/ask";
 import { MappingFileRecord} from "./types";
 import { generateAutoMappings, generateRepairCode } from "../../../src/features/ai/service/datamapper/datamapper";
 import { ArtifactNotificationHandler, ArtifactsUpdated } from "../../utils/project-artifacts-handler";
+import { CopilotEventHandler } from "../../../src/features/ai/service/event";
 
 // const BACKEND_BASE_URL = BACKEND_URL.replace(/\/v2\.0$/, "");
 //TODO: Temp workaround as custom domain seem to block file uploads
@@ -154,15 +155,18 @@ async function convertAttachmentToFileData(attachment: Attachment): Promise<File
 // Processes data mapper model and optional mapping instruction files to generate mapping expressions
 export async function generateMappingExpressionsFromModel(
     dataMapperModel: DMModel,
-    mappingInstructionFiles: Attachment[] = []
+    mappingInstructionFiles: Attachment[] = [],
+    eventHandler: CopilotEventHandler
 ): Promise<Mapping[]> {
     let dataMapperResponse: DataMapperModelResponse = {
         mappingsModel: dataMapperModel as DMModel
     };
     if (mappingInstructionFiles.length > 0) {
+        eventHandler({ type: "content_block", content: "\n<progress>Processing mapping hints from attachments...</progress>" });
         const enhancedResponse = await enrichModelWithMappingInstructions(mappingInstructionFiles, dataMapperResponse);
         dataMapperResponse = enhancedResponse as DataMapperModelResponse;
     }
+    eventHandler({ type: "content_block", content: "\n<progress>Generating data mappings...</progress>" });
 
     const generatedMappings = await generateAutoMappings(dataMapperResponse);
     return generatedMappings.map(mapping => ({