fix: Migrates gemini cache from os temp dir to workspace cache

vscode/src/extension.ts

@@ -3,7 +3,6 @@
 // Import the module and reference it with the alias vscode in your code below
 import * as vscode from 'vscode';
 
-import { createInlineCodeCompletion } from './tools/create/inline_code_completion';
 import { makeHttpRequest } from './repository/http-utils';
 import { activateTelemetry, logEvent } from './utilities/telemetry-reporter';
 import * as dotenv from 'dotenv';

vscode/src/repository/gemini-repository.ts

@@ -8,6 +8,8 @@ import { getReferenceEditor } from "../utilities/state-objects";
 import { GenerationRepository } from "./generation-repository";
 import { extractReferenceTextFromEditor } from "../utilities/code-processing";
 import { logError } from "../utilities/telemetry-reporter";
+import { CacheManager } from "../utilities/cache-manager";
+import { computeCodehash } from "../shared/utils";
 
 function handleError(error: Error, userFriendlyMessage: string): never {
     console.error(error);
@@ -22,12 +24,37 @@ export class GeminiRepository extends GenerationRepository {
     constructor(apiKey: string) {
         super(apiKey);
         this.genAI = new GoogleGenerativeAI(apiKey);
-        this.ensureCacheDirExists().catch(error => {
-            handleError(error, 'Failed to initialize the cache directory.');
-        });
+        this.migrateCache();
         GeminiRepository._instance = this;
     }
 
+    /**
+     * Migrates the cache from the operating system's temporary directory to the extension's cache directory.
+     * 
+     * The cache was previously stored in the operating system's temporary directory, but it is now being moved to the extension's cache directory.
+     * This migration ensures that the cache is stored in a secure location and avoids potential issues with file permissions.
+     * 
+     * This is a one-time operation and will be removed in future releases.
+     */
+    private async migrateCache() {
+        try {
+            // TODO: Need to update the migration logic for the new cache structure
+            // Check if the cache file exists in the old location
+            if (fs.existsSync(this.cacheFilePath)) {
+                // Move the cache file to the CacheManager
+                const cacheData = await fs.promises.readFile(this.cacheFilePath, 'utf8');
+                console.log(cacheData);
+                if (cacheData.length > 0) {
+                    await CacheManager.getInstance().setGeminiCache(this.codehashCache);
+                }
+                // Delete the old cache file
+                await fs.promises.unlink(this.cacheFilePath);
+            }
+        } catch (error) {
+            console.error('Failed to migrate the cache:', error);
+        }
+    }
+
     public static getInstance(): GeminiRepository {
         return GeminiRepository._instance;
     }
@@ -124,7 +151,7 @@ export class GeminiRepository extends GenerationRepository {
             throw new Error('No workspace folders found.');
         }
         const projectFolder = workspaceFolders[0].uri.fsPath; // Assuming single root workspace
-        const hash = this.computeCodehash(projectFolder); // Hash the path for uniqueness
+        const hash = computeCodehash(projectFolder); // Hash the path for uniqueness
         // Use os.tmpdir() to get the system's temporary directory
         const tempDir = require('os').tmpdir();
         return require('path').join(tempDir, 'flutterGPT', `${hash}.codehashCache.json`);
@@ -135,8 +162,7 @@ export class GeminiRepository extends GenerationRepository {
     private async saveCache() {
         try {
             const cacheData = JSON.stringify(this.codehashCache);
-            const cachePath = this.cacheFilePath;
-            await fs.promises.writeFile(cachePath, cacheData, { encoding: 'utf8', mode: 0o600 }); // Sets the file mode to read/write for the owner only
+            await CacheManager.getInstance().setGeminiCache(this.codehashCache);
         } catch (error) {
             if (error instanceof Error) {
                 handleError(error, 'Failed to save the cache data.');
@@ -148,37 +174,18 @@ export class GeminiRepository extends GenerationRepository {
 
     private async loadCache() {
         try {
-            if (fs.existsSync(this.cacheFilePath)) {
-                const cacheData = await fs.promises.readFile(this.cacheFilePath, 'utf8');
+            const cacheData = await CacheManager.getInstance().getGeminiCache();
+            if (cacheData) {
                 this.codehashCache = JSON.parse(cacheData);
             }
+            else {
+                this.codehashCache = {};
+            }
         } catch (error) {
             console.error("Error loading cache: ", error);
         }
     }
 
-    // Ensure the directory exists and has the correct permissions
-    private async ensureCacheDirExists() {
-        const cacheDir = path.dirname(this.cacheFilePath);
-        try {
-            // Create the directory if it doesn't exist
-            if (!fs.existsSync(cacheDir)) {
-                await fs.promises.mkdir(cacheDir, { recursive: true, mode: 0o700 }); // Sets the directory mode to read/write/execute for the owner only
-            }
-        } catch (error: any) {
-            if (error.code !== 'EEXIST') {
-                handleError(error as Error, 'Failed to create a secure cache directory.');
-            }
-        }
-    }
-
-    // Compute a codehash for file contents
-    private computeCodehash(fileContents: string): string {
-        // Normalize the file content by removing whitespace and newlines
-        const normalizedContent = fileContents.replace(/\s+/g, '');
-        return crypto.createHash('sha256').update(normalizedContent).digest('hex');
-    }
-
     // Find 5 closest dart files for query
     public async findClosestDartFiles(query: string, view?: vscode.WebviewView, shortcut: boolean = false, filepath: string = ''): Promise<string> {
         //start timer
@@ -219,7 +226,7 @@ export class GeminiRepository extends GenerationRepository {
                     const document = await vscode.workspace.openTextDocument(file);
                     const relativePath = vscode.workspace.asRelativePath(file, false);
                     const text = `File name: ${file.path.split('/').pop()}\nFile path: ${relativePath}\nFile code:\n\n\`\`\`dart\n${document.getText()}\`\`\`\n\n------\n\n`;
-                    const codehash = this.computeCodehash(text);
+                    const codehash = computeCodehash(text);
                     return {
                         text,
                         path: file.path,

vscode/src/shared/utils.ts

@@ -6,6 +6,7 @@ import { GeminiRepository } from '../repository/gemini-repository';
 import { ILspAnalyzer } from './types/LspAnalyzer';
 import { Outline } from './types/custom_protocols';
 import { SemanticTokensRegistrationType, SemanticTokensProviderShape, SymbolKind } from 'vscode-languageclient';
+import * as crypto from 'crypto';
 
 export async function getCodeForElementAtRange(analyzer: ILspAnalyzer, document: vscode.TextDocument, range: vscode.Range): Promise<string | undefined> {
 	const outline = (await analyzer.fileTracker.waitForOutline(document));
@@ -46,7 +47,7 @@ export async function getCodeForElementAtRange(analyzer: ILspAnalyzer, document:
 }
 
 export async function cursorIsAt(type: String, analyzer: ILspAnalyzer, document: vscode.TextDocument, activeTextEditor: vscode.TextEditor | undefined, range: vscode.Range, strict: boolean = true): Promise<{ symbolRange: vscode.Range, symbol: Outline } | undefined> {
-	if (!activeTextEditor){
+	if (!activeTextEditor) {
 		return undefined;
 	}
 	const position = activeTextEditor.selection.active;
@@ -177,3 +178,9 @@ export async function getCodeForRange(uri: vscode.Uri, range: vscode.Range): Pro
 
 	return undefined;
 }
+
+export function computeCodehash(fileContents: string): string {
+	// Normalize the file content by removing whitespace and newlines
+	const normalizedContent = fileContents.replace(/\s+/g, '');
+	return crypto.createHash('sha256').update(normalizedContent).digest('hex');
+}

vscode/src/tools/create/inline_code_completion.ts

@@ -68,7 +68,7 @@ function replaceLineOfCode(line: number, replaceString: string) {
     });
 }
 
-export async function createInlineCodeCompletion(geminiRepo: GeminiRepository) {
+export async function createInlineCodeCompletion() {
     // manual trigger using shortcut ctrl+space
     logEvent('create-inline-code-completion', { 'type': "create" });
     vscode.window.withProgress({
@@ -100,6 +100,7 @@ async function generateSuggestions(): Promise<string[]> {
         const fileContent = editor.document.getText();
         const filepath = editor.document.fileName; // don't include currentFile in most relevant files.
         console.log("Current file path:", filepath);
+        ///TODO: Replace with Generation Repository
         var relevantFiles = await GeminiRepository.getInstance().findClosestDartFiles("Current file content:" + editor.document.getText() + "\n\n" + "Line of code:" + currentLineContent, undefined, true, filepath);
         // Add code for all the elements used in the file.
         const contextualCode = await new ContextualCodeProvider().getContextualCodeForCompletion(editor.document, getDartAnalyser());

vscode/src/utilities/cache-manager.ts

@@ -1,6 +1,8 @@
 import * as vscode from 'vscode';
 import { logError } from './telemetry-reporter';
-
+import { ContentEmbedding } from '@google/generative-ai';
+import * as path from 'path';
+import { computeCodehash } from '../shared/utils';
 export class CacheManager {
     private static instance: CacheManager;
     private globalState: vscode.Memento;
@@ -58,4 +60,77 @@ export class CacheManager {
             return 0;
         }
     }
+
+    async setGeminiCache(cacheData: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } }): Promise<void> {
+        const excludePatterns = "**/{android,ios,web,linux,macos,windows,.dart_tool}/**";
+        const pubspecs = await vscode.workspace.findFiles("**/pubspec.yaml", excludePatterns);
+        if (pubspecs.length === 0) {
+            throw new Error("No pubspec.yaml found in the workspace.");
+        }
+        // Get all the flutter projects in the workspace
+        const flutterProjects = pubspecs.map((uri) => uri.fsPath);
+        const cache: { [flutterProject: string]: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } } } = {};
+
+        // ITERATE OVER ALL THE FILES IN THE CACHE
+        // Find the flutter project for the file
+        // Add the file to the cache of that flutter project
+        for (const filePath in cacheData) {
+            const parentProjectPath = this.findParentFlutterProject(filePath, flutterProjects);
+            if (parentProjectPath) {
+                const key = "gemini-cache-" + computeCodehash(parentProjectPath);
+                await this.setGlobalValue(key, JSON.stringify(cacheData));
+            }
+        }
+
+    }
+
+    async getGeminiCache(): Promise<string | undefined> {
+        const excludePatterns = "**/{android,ios,web,linux,macos,windows,.dart_tool}/**";
+        const pubspecs = await vscode.workspace.findFiles("**/pubspec.yaml", excludePatterns);
+        if (pubspecs.length === 0) {
+            throw new Error("No pubspec.yaml found in the workspace.");
+        }
+
+        // Get all the flutter projects in the workspace
+        const flutterProjects = pubspecs.map((uri) => uri.fsPath);
+
+        const activeCache: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } } = {};
+        // Return cache only for the parent flutter project of the current workspace
+        for (const projectPath of flutterProjects) {
+            const projectDir = path.dirname(projectPath);
+            const key = "gemini-cache-" + projectDir;
+            // await this.setGlobalValue<String>(key, "value");
+            const cacheString = await this.getGlobalValue<string>(key);
+            if (!cacheString) {
+                continue;
+            }
+            const cache = JSON.parse(cacheString);
+            // Add the cache for the project to the activeCache
+            Object.assign(activeCache, cache);
+        }
+
+        return JSON.stringify(activeCache);
+    }
+
+    // Helper function to find parent Flutter project
+    private findParentFlutterProject(filePath: string, flutterProjects: string[]) {
+        let parentProjectPath = null;
+        let maxCommonLength = -1;
+
+        for (const projectPath of flutterProjects) {
+            const projectDir = path.dirname(projectPath);
+
+            // Check if the current projectDir is a prefix of the filePath
+            if (filePath.startsWith(projectDir)) {
+                const commonLength = projectDir.length;
+
+                if (commonLength > maxCommonLength) {
+                    maxCommonLength = commonLength;
+                    parentProjectPath = projectDir;
+                }
+            }
+        }
+
+        return parentProjectPath;
+    }
 }

vscode/src/utilities/command-manager.ts

@@ -69,7 +69,7 @@ export function registerCommand(
 }
 
 export function initCommands(context: vscode.ExtensionContext, geminiRepo: any, analyzer: any, flutterGPTViewProvider: FlutterGPTViewProvider) {
-    const generationRepository: GenerationRepository = getUserPrefferedModel();
+    const generationRepository: GenerationRepository =  geminiRepo; //TODO: Use up in the tree getUserPrefferedModel();
     // List of commands to register, with their names and options.
     const commands = [
         { name: 'dashai.attachToDash', handler: () => addToReference(context.globalState, flutterGPTViewProvider), options: { isCommand: true, isMenu: true, isShortcut: false } },
@@ -81,6 +81,7 @@ export function initCommands(context: vscode.ExtensionContext, geminiRepo: any,
         { name: 'dashai.optimizeCode', handler: (aiRepo: GenerationRepository, globalState: vscode.Memento, range: vscode.Range, anlyzer: ILspAnalyzer, elementName: string | undefined) => optimizeCode(generationRepository, context.globalState, range, anlyzer, elementName, context), options: { isCommand: true, isMenu: false, isShortcut: false } },
         { name: 'dashai.createInlineCodeCompletion', handler: () => createInlineCodeCompletion(geminiRepo), options: { isCommand: true, isMenu: true, isShortcut: true } },
         { name: 'dashai.clearChat', handler: () => flutterGPTViewProvider?.postMessageToWebview({ type: 'clearCommandDeck' }), options: { isCommand: true, isMenu: false, isShortcut: false } }
+
         // Add more commands as needed.
     ];