Use VS Code native file dialog for AI panel attachments

Author: VellummyilumVinoth

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

@@ -56,6 +56,8 @@ import {
     RunningServiceInfo,
     StopRunningServiceRequest,
     RunServiceRequest,
+    SelectContextFilesRequest,
+    Attachment,
 } from "./interfaces";
 
 export interface AIPanelAPI {
@@ -140,4 +142,8 @@ export interface AIPanelAPI {
     // Vertex AI BYOK Helpers
     // ==================================
     getDefaultVertexCredsPath: () => Promise<string>;
+    // ==================================
+    // File Attachment via VS Code Dialog
+    // ==================================
+    selectContextFiles: (params: SelectContextFilesRequest) => Promise<Attachment[]>;
 }

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

@@ -240,6 +240,10 @@ export interface ExistingFunction {
 // ==================================
 // Attachment-Related Interfaces
 // ==================================
+export interface SelectContextFilesRequest {
+    command: Command | null;
+}
+
 export interface Attachment {
     name: string;
     path?: string

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

@@ -58,6 +58,8 @@ import {
     RunningServiceInfo,
     StopRunningServiceRequest,
     RunServiceRequest,
+    SelectContextFilesRequest,
+    Attachment,
 } from "./interfaces";
 import { RequestType, NotificationType } from "vscode-messenger-common";
 
@@ -121,3 +123,4 @@ export const stopRunningService: RequestType<StopRunningServiceRequest, boolean>
 export const runService: RequestType<RunServiceRequest, boolean> = { method: `${_preFix}/runService` };
 export const runningServicesChanged: NotificationType<RunningServiceInfo[]> = { method: `${_preFix}/runningServicesChanged` };
 export const getDefaultVertexCredsPath: RequestType<void, string> = { method: `${_preFix}/getDefaultVertexCredsPath` };
+export const selectContextFiles: RequestType<SelectContextFilesRequest, Attachment[]> = { method: `${_preFix}/selectContextFiles` };

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

@@ -106,6 +106,8 @@ import {
     runService,
     RunServiceRequest,
     getDefaultVertexCredsPath,
+    selectContextFiles,
+    SelectContextFilesRequest,
 } from "@wso2/ballerina-core";
 import { workspace } from 'vscode';
 import { Messenger } from "vscode-messenger";
@@ -182,6 +184,7 @@ export function registerAiPanelRpcHandlers(messenger: Messenger) {
     messenger.onRequest(stopRunningService, (args: StopRunningServiceRequest) => rpcManger.stopRunningService(args));
     messenger.onRequest(runService, (args: RunServiceRequest) => rpcManger.runService(args));
     messenger.onRequest(getDefaultVertexCredsPath, () => rpcManger.getDefaultVertexCredsPath());
+    messenger.onRequest(selectContextFiles, (args: SelectContextFilesRequest) => rpcManger.selectContextFiles(args));
 
     // Push updates to the webview whenever the set of running services changes.
     runningServicesManager.onChange = (services) => {

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

@@ -24,6 +24,8 @@ import {
     AIPanelPrompt,
     AbortAIGenerationRequest,
     AddFilesToProjectRequest,
+    Attachment,
+    AttachmentStatus,
     CheckpointInfo,
     Command,
     DocGenerationRequest,
@@ -42,6 +44,7 @@ import {
     RestoreCheckpointRequest,
     SemanticDiffRequest,
     SemanticDiffResponse,
+    SelectContextFilesRequest,
     SubmitFeedbackRequest,
     TestGenerationMentions,
     UIChatMessage,
@@ -864,4 +867,39 @@ User reverted the last made changes. The files have been restored to the state b
         return "";
     }
 
+    async selectContextFiles(params: SelectContextFilesRequest): Promise<Attachment[]> {
+        const projectPath = StateMachine.context().projectPath;
+        const defaultUri = projectPath ? vscode.Uri.file(projectPath) : vscode.Uri.file(os.homedir());
+
+        const uris = await vscode.window.showOpenDialog({
+            canSelectMany: true,
+            canSelectFiles: true,
+            canSelectFolders: false,
+            defaultUri,
+            openLabel: "Attach",
+        });
+
+        if (!uris || uris.length === 0) {
+            return [];
+        }
+
+        const MAX_FILE_SIZE = 5 * 1024 * 1024;
+        const isDataMapper = params.command === Command.DataMap || params.command === Command.TypeCreator;
+
+        return uris.map((uri) => {
+            const filePath = uri.fsPath;
+            const name = path.basename(filePath);
+            try {
+                if (fs.statSync(filePath).size > MAX_FILE_SIZE) {
+                    return { name, status: AttachmentStatus.FileSizeExceeded };
+                }
+                const content = isDataMapper
+                    ? fs.readFileSync(filePath).toString("base64")
+                    : fs.readFileSync(filePath, "utf-8");
+                return { name, content, status: AttachmentStatus.Success };
+            } catch {
+                return { name, status: AttachmentStatus.UnknownError };
+            }
+        });
+    }
 }

workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts

@@ -117,6 +117,9 @@ import {
     RunServiceRequest,
     runService,
     getDefaultVertexCredsPath,
+    selectContextFiles,
+    SelectContextFilesRequest,
+    Attachment,
 } from "@wso2/ballerina-core";
 import { HOST_EXTENSION } from "vscode-messenger-common";
 import { Messenger } from "vscode-messenger-webview";
@@ -361,4 +364,8 @@ export class AiPanelRpcClient implements AIPanelAPI {
     getDefaultVertexCredsPath(): Promise<string> {
         return this._messenger.sendRequest(getDefaultVertexCredsPath, HOST_EXTENSION);
     }
+
+    selectContextFiles(params: SelectContextFilesRequest): Promise<Attachment[]> {
+        return this._messenger.sendRequest(selectContextFiles, HOST_EXTENSION, params);
+    }
 }

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx

@@ -2180,6 +2180,8 @@ const AIChat: React.FC = () => {
                                 multiple: true,
                                 acceptResolver: acceptResolver,
                                 handleAttachmentSelection: handleAttachmentSelection,
+                                onAttachClick: (command) =>
+                                    rpcClient.getAiPanelRpcClient().selectContextFiles({ command }),
                             }}
                             suggestedCommandTemplates={footerSuggestedCommandTemplates}
                             inputPlaceholder={footerInputPlaceholder}

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChatInput/hooks/useAttachments.ts

@@ -17,12 +17,15 @@
  */
 
 import { useState, useRef, ChangeEvent } from "react";
-import { Attachment, Command } from "@wso2/ballerina-core";
+import { Attachment, AttachmentStatus, Command } from "@wso2/ballerina-core";
 
 export interface AttachmentOptions {
     multiple: boolean;
     acceptResolver: (command: Command | null) => string;
     handleAttachmentSelection: (e: ChangeEvent<HTMLInputElement>, command: Command | null) => Promise<Attachment[]>;
+    /** If set, clicking "Attach" uses this instead of the hidden file input,
+    * allowing the VS Code file dialog to open (instead of the browser file picker). */
+    onAttachClick?: (command: Command | null) => Promise<Attachment[]>;
 }
 
 interface UseAttachmentsProps {
@@ -34,9 +37,26 @@ export function useAttachments({ attachmentOptions, activeCommand }: UseAttachme
     const [attachments, setAttachments] = useState<Attachment[]>([]);
     const fileInputRef = useRef<HTMLInputElement | null>(null);
 
-    // open file input
-    function handleAttachClick() {
-        if (fileInputRef.current) {
+    // open file picker
+    async function handleAttachClick() {
+        if (attachmentOptions.onAttachClick) {
+            const results = await attachmentOptions.onAttachClick(activeCommand);
+            setAttachments((prev) => {
+                const updated = [...prev];
+                results
+                    .filter((newFile) => newFile.status === AttachmentStatus.Success)
+                    .forEach((newFile) => {
+                        const existingIndex = updated.findIndex(
+                            (existing) => existing.name === newFile.name && existing.content === newFile.content
+                        );
+                        if (existingIndex !== -1) {
+                            updated.splice(existingIndex, 1);
+                        }
+                        updated.push(newFile);
+                    });
+                return updated;
+            });
+        } else if (fileInputRef.current) {
             fileInputRef.current.click();
         }
     }