fix(featureDev): Prevent crash on large repos or unsupported encoding during file collection

Author: jsalles

packages/core/src/amazonqFeatureDev/constants.ts

@@ -14,6 +14,9 @@ export const featureDevChat = 'featureDevChat'
 
 export const featureName = 'Amazon Q feature development'
 
+// Max allowed size for file collection
+export const maxRepoSize = 200 * 1024 * 1024
+
 // License text that's used in codewhisperer reference log
 export const referenceLogText = (reference: CodeReference) =>
     `[${new Date().toLocaleString()}] Accepted recommendation from Amazon Q. Code provided with reference under <a href="${LicenseUtil.getLicenseHtml(

packages/core/src/amazonqFeatureDev/util/files.ts

@@ -12,7 +12,7 @@ import { Uri } from 'vscode'
 import { GitIgnoreFilter } from './gitignore'
 
 import AdmZip from 'adm-zip'
-import { PrepareRepoFailedError } from '../errors'
+import { ContentLengthError, PrepareRepoFailedError } from '../errors'
 import { getLogger } from '../../shared/logger/logger'
 import { maxFileSizeBytes } from '../limits'
 import { createHash } from 'crypto'
@@ -21,6 +21,7 @@ import { ToolkitError } from '../../shared/errors'
 import { AmazonqCreateUpload, Metric } from '../../shared/telemetry/telemetry'
 import { TelemetryHelper } from './telemetryHelper'
 import { sanitizeFilename } from '../../shared/utilities/textUtilities'
+import { maxRepoSize } from '../constants'
 
 export function getExcludePattern(additionalPatterns: string[] = []) {
     const globAlwaysExcludedDirs = getGlobDirExcludedPatterns().map(pattern => `**/${pattern}/*`)
@@ -94,6 +95,7 @@ export async function collectFiles(
         return prefix === '' ? path : `${prefix}/${path}`
     }
 
+    let totalSize = 0
     for (const rootPath of sourcePaths) {
         const allFiles = await vscode.workspace.findFiles(
             new vscode.RelativePattern(rootPath, '**'),
@@ -102,29 +104,43 @@ export async function collectFiles(
         const files = respectGitIgnore ? await filterOutGitignoredFiles(rootPath, allFiles) : allFiles
 
         for (const file of files) {
-            try {
-                const fileContent = await SystemUtilities.readFile(file, new TextDecoder('utf8', { fatal: true }))
-                const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
+            const fileStat = await vscode.workspace.fs.stat(file)
+            const { relativePath, fileContent } = await readFile(file, workspaceFolders)
+            if (relativePath === undefined || fileContent === undefined) {
+                continue
+            }
 
-                if (relativePath) {
-                    storage.push({
-                        workspaceFolder: relativePath.workspaceFolder,
-                        relativeFilePath: relativePath.relativePath,
-                        fileUri: file,
-                        fileContent: fileContent,
-                        zipFilePath: prefixWithFolderPrefix(relativePath.workspaceFolder, relativePath.relativePath),
-                    })
-                }
-            } catch (error) {
-                getLogger().debug(
-                    `featureDev: Failed to read file ${file.fsPath} when collecting repository: ${error}. Skipping the file`
-                )
+            totalSize += fileStat.size
+            if (totalSize > maxRepoSize) {
+                throw new ContentLengthError()
             }
+
+            storage.push({
+                workspaceFolder: relativePath.workspaceFolder,
+                relativeFilePath: relativePath.relativePath,
+                fileUri: file,
+                fileContent: fileContent,
+                zipFilePath: prefixWithFolderPrefix(relativePath.workspaceFolder, relativePath.relativePath),
+            })
         }
     }
     return storage
 }
 
+const readFile = async (file: vscode.Uri, workspaceFolders: CurrentWsFolders) => {
+    try {
+        const fileContent = await SystemUtilities.readFile(file, new TextDecoder('utf8', { fatal: false }))
+        const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
+        return { fileContent, relativePath }
+    } catch (error) {
+        getLogger().debug(
+            `featureDev: Failed to read file ${file.fsPath} when collecting repository. Skipping the file`
+        )
+    }
+
+    return { fileContent: undefined, relativePath: undefined }
+}
+
 const getSha256 = (file: Buffer) => createHash('sha256').update(file).digest('base64')
 
 /**
@@ -137,9 +153,9 @@ export async function prepareRepoData(
     span: Metric<AmazonqCreateUpload>
 ) {
     try {
+        const files = await collectFiles(repoRootPaths, workspaceFolders, true)
         const zip = new AdmZip()
 
-        const files = await collectFiles(repoRootPaths, workspaceFolders, true)
         let totalBytes = 0
         for (const file of files) {
             const fileSize = (await vscode.workspace.fs.stat(file.fileUri)).size
@@ -163,6 +179,9 @@ export async function prepareRepoData(
         }
     } catch (error) {
         getLogger().debug(`featureDev: Failed to prepare repo: ${error}`)
+        if (error instanceof ContentLengthError) {
+            throw error
+        }
         throw new PrepareRepoFailedError()
     }
 }