feat(amazonq): /doc: add support for infrastructure diagrams

Author: Viktor Shesternyak

packages/amazonq/.changes/next-release/Feature-088c8b44-eb51-4495-84fb-2e02878ff4bc.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Amazon Q /doc: Add support for infrastructure diagrams"
+}

packages/core/src/amazonq/session/sessionState.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import { ToolkitError } from '../../shared/errors'
 import globals from '../../shared/extensionGlobals'
 import { getLogger } from '../../shared/logger/logger'
-import { telemetry } from '../../shared/telemetry/telemetry'
+import { AmazonqCreateUpload, Span, telemetry } from '../../shared/telemetry/telemetry'
 import { VirtualFileSystem } from '../../shared/virtualFilesystem'
 import { CodeReference, UploadHistory } from '../webview/ui/connector'
 import { AuthUtil } from '../../codewhisperer/util/authUtil'
@@ -27,6 +27,7 @@ import {
 } from '../commons/types'
 import { prepareRepoData, getDeletedFileInfos, registerNewFiles } from '../util/files'
 import { uploadCode } from '../util/upload'
+import { TelemetryHelper } from '../util/telemetryHelper'
 
 export const EmptyCodeGenID = 'EMPTY_CURRENT_CODE_GENERATION_ID'
 
@@ -227,7 +228,7 @@ export abstract class BasePrepareCodeGenState implements SessionState {
                 amazonqConversationId: this.config.conversationId,
                 credentialStartUrl: AuthUtil.instance.startUrl,
             })
-            const { zipFileBuffer, zipFileChecksum } = await prepareRepoData(
+            const { zipFileBuffer, zipFileChecksum } = await this.prepareProjectZip(
                 this.config.workspaceRoots,
                 this.config.workspaceFolders,
                 action.telemetry,
@@ -251,6 +252,15 @@ export abstract class BasePrepareCodeGenState implements SessionState {
         const nextState = this.createNextState({ ...this.config, uploadId })
         return nextState.interact(action)
     }
+
+    protected async prepareProjectZip(
+        workspaceRoots: string[],
+        workspaceFolders: CurrentWsFolders,
+        telemetry: TelemetryHelper,
+        span: Span<AmazonqCreateUpload>
+    ) {
+        return await prepareRepoData(workspaceRoots, workspaceFolders, telemetry, span)
+    }
 }
 
 export interface CodeGenerationParams {

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

@@ -5,7 +5,12 @@
 
 import * as vscode from 'vscode'
 import * as path from 'path'
-import { collectFiles, getWorkspaceFoldersByPrefixes } from '../../shared/utilities/workspaceUtils'
+import {
+    collectFiles,
+    CollectFilesFileFilter,
+    DefaultExtraExcludePatterns,
+    getWorkspaceFoldersByPrefixes,
+} from '../../shared/utilities/workspaceUtils'
 
 import { ContentLengthError, PrepareRepoFailedError } from '../../amazonqFeatureDev/errors'
 import { getLogger } from '../../shared/logger/logger'
@@ -30,6 +35,13 @@ export async function checkForDevFile(root: string) {
     return hasDevFile
 }
 
+function isInfraDiagramFile(relativePath: string) {
+    return (
+        relativePath.toLowerCase().endsWith(path.join('docs', 'infra.dot')) ||
+        relativePath.toLowerCase().endsWith(path.join('docs', 'infra.svg'))
+    )
+}
+
 /**
  * given the root path of the repo it zips its files in memory and generates a checksum for it.
  */
@@ -38,13 +50,38 @@ export async function prepareRepoData(
     workspaceFolders: CurrentWsFolders,
     telemetry: TelemetryHelper,
     span: Span<AmazonqCreateUpload>,
-    zip: ZipStream = new ZipStream()
+    zip: ZipStream = new ZipStream(),
+    isIncludeInfraDiagram: boolean = false
 ) {
     try {
         const autoBuildSetting = CodeWhispererSettings.instance.getAutoBuildSetting()
         const useAutoBuildFeature = autoBuildSetting[repoRootPaths[0]] ?? false
+        const extraExcludeFilePatterns: string[] = []
+        let extraFileFilterFn: CollectFilesFileFilter | undefined = undefined
+
         // We only respect gitignore file rules if useAutoBuildFeature is on, this is to avoid dropping necessary files for building the code (e.g. png files imported in js code)
-        const files = await collectFiles(repoRootPaths, workspaceFolders, true, maxRepoSizeBytes, !useAutoBuildFeature)
+        if (!useAutoBuildFeature) {
+            if (isIncludeInfraDiagram) {
+                // ensure svg is not filtered out by files search
+                extraExcludeFilePatterns.push(...DefaultExtraExcludePatterns.filter((p) => !p.endsWith('.svg')))
+                // ensure only infra diagram is included from all svg files
+                extraFileFilterFn = (relativePath: string) => {
+                    if (!relativePath.toLowerCase().endsWith('.svg')) {
+                        return false
+                    }
+                    return !isInfraDiagramFile(relativePath)
+                }
+            } else {
+                extraExcludeFilePatterns.push(...DefaultExtraExcludePatterns)
+            }
+        }
+
+        const files = await collectFiles(repoRootPaths, workspaceFolders, {
+            maxSizeLimitInBytes: maxRepoSizeBytes,
+            useGitIgnoreFileAsFilter: true,
+            extraExcludeFilePatterns,
+            extraFileFilterFn,
+        })
 
         let totalBytes = 0
         const ignoredExtensionMap = new Map<string, number>()
@@ -62,9 +99,15 @@ export async function prepareRepoData(
             }
             const isCodeFile_ = isCodeFile(file.relativeFilePath)
             const isDevFile = file.relativeFilePath === 'devfile.yaml'
-            // When useAutoBuildFeature is on, only respect the gitignore rules filtered earlier and apply the size limit, otherwise, exclude all non code files and gitignore files
-            const isNonCodeFileAndIgnored = useAutoBuildFeature ? false : !isCodeFile_ || isDevFile
-            if (fileSize >= maxFileSizeBytes || isNonCodeFileAndIgnored) {
+            const isInfraDiagramFileExt = isInfraDiagramFile(file.relativeFilePath)
+
+            let isExcludeFile = fileSize >= maxFileSizeBytes
+            // When useAutoBuildFeature is on, only respect the gitignore rules filtered earlier and apply the size limit
+            if (!isExcludeFile && !useAutoBuildFeature) {
+                isExcludeFile = isDevFile || (!isCodeFile_ && (!isIncludeInfraDiagram || !isInfraDiagramFileExt))
+            }
+
+            if (isExcludeFile) {
                 if (!isCodeFile_) {
                     const re = /(?:\.([^.]+))?$/
                     const extensionArray = re.exec(file.relativeFilePath)
@@ -77,6 +120,7 @@ export async function prepareRepoData(
                 }
                 continue
             }
+
             totalBytes += fileSize
             // Paths in zip should be POSIX compliant regardless of OS
             // Reference: https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT

packages/core/src/amazonqDoc/controllers/chat/controller.ts

@@ -35,7 +35,7 @@ import { BaseChatSessionStorage } from '../../../amazonq/commons/baseChatStorage
 import { DocMessenger } from '../../messenger'
 import { AuthController } from '../../../amazonq/auth/controller'
 import { openUrl } from '../../../shared/utilities/vsCodeUtils'
-import { openDeletedDiff, openDiff } from '../../../amazonq/commons/diff'
+import { createAmazonQUri, openDeletedDiff, openDiff } from '../../../amazonq/commons/diff'
 import {
     getWorkspaceFoldersByPrefixes,
     getWorkspaceRelativePath,
@@ -204,7 +204,13 @@ export class DocController {
                     uploadId = session?.state?.uploadHistory[codeGenerationId].uploadId
                 }
                 const rightPath = path.join(uploadId, zipFilePath)
-                await openDiff(pathInfos.absolutePath, rightPath, tabId, this.scheme)
+                if (rightPath.toLowerCase().endsWith('.svg')) {
+                    // use open instead of diff for svg
+                    const rightPathUri = createAmazonQUri(rightPath, tabId, this.scheme)
+                    await vscode.commands.executeCommand('vscode.open', rightPathUri)
+                } else {
+                    await openDiff(pathInfos.absolutePath, rightPath, tabId, this.scheme)
+                }
             }
         }
     }

packages/core/src/amazonqDoc/session/sessionState.ts

@@ -8,7 +8,7 @@ import { DocGenerationStep, docScheme, getFileSummaryPercentage, Mode } from '..
 
 import { i18n } from '../../shared/i18n-helper'
 
-import { NewFileInfo, SessionState, SessionStateAction, SessionStateConfig } from '../types'
+import { CurrentWsFolders, NewFileInfo, SessionState, SessionStateAction, SessionStateConfig } from '../types'
 import {
     ContentLengthError,
     DocServiceError,
@@ -23,6 +23,10 @@ import {
 import { DocMessenger } from '../messenger'
 import { BaseCodeGenState, BasePrepareCodeGenState, CreateNextStateParams } from '../../amazonq/session/sessionState'
 import { Intent } from '../../amazonq/commons/types'
+import { TelemetryHelper } from '../../amazonq/util/telemetryHelper'
+import { AmazonqCreateUpload, Span } from '../../shared/telemetry/telemetry'
+import { prepareRepoData } from '../../amazonq/util/files'
+import { ZipStream } from '../../shared/utilities/zipStream'
 
 export class DocCodeGenState extends BaseCodeGenState {
     protected handleProgress(messenger: DocMessenger, action: SessionStateAction, detail?: string): void {
@@ -132,4 +136,13 @@ export class DocPrepareCodeGenState extends BasePrepareCodeGenState {
     protected override createNextState(config: SessionStateConfig): SessionState {
         return super.createNextState(config, DocCodeGenState)
     }
+
+    protected override async prepareProjectZip(
+        workspaceRoots: string[],
+        workspaceFolders: CurrentWsFolders,
+        telemetry: TelemetryHelper,
+        span: Span<AmazonqCreateUpload>
+    ) {
+        return await prepareRepoData(workspaceRoots, workspaceFolders, telemetry, span, new ZipStream(), true)
+    }
 }

packages/core/src/amazonqFeatureDev/session/sessionState.ts

@@ -77,7 +77,9 @@ export class MockCodeGenState implements SessionState {
             const files = await collectFiles(
                 this.config.workspaceFolders.map((f) => path.join(f.uri.fsPath, './mock-data')),
                 this.config.workspaceFolders,
-                false
+                {
+                    useGitIgnoreFileAsFilter: false,
+                }
             )
             const newFileContents = files.map((f) => ({
                 zipFilePath: f.zipFilePath,

packages/core/src/codewhisperer/util/zipUtil.ts

@@ -410,12 +410,9 @@ export class ZipUtil {
             return
         }
 
-        const sourceFiles = await collectFiles(
-            projectPaths,
-            vscode.workspace.workspaceFolders as CurrentWsFolders,
-            true,
-            this.getProjectScanPayloadSizeLimitInBytes()
-        )
+        const sourceFiles = await collectFiles(projectPaths, vscode.workspace.workspaceFolders as CurrentWsFolders, {
+            maxSizeLimitInBytes: this.getProjectScanPayloadSizeLimitInBytes(),
+        })
         for (const file of sourceFiles) {
             const projectName = path.basename(file.workspaceFolder.uri.fsPath)
             const zipEntryPath = this.getZipEntryPath(projectName, file.relativeFilePath)

packages/core/src/shared/utilities/workspaceUtils.ts

@@ -19,6 +19,7 @@ import * as parser from '@gerhobbelt/gitignore-parser'
 import fs from '../fs/fs'
 import { ChildProcess } from './processUtils'
 import { isWin } from '../vscode/env'
+import { maxRepoSizeBytes } from '../../amazonqFeatureDev/constants'
 
 type GitIgnoreRelativeAcceptor = {
     folderPath: string
@@ -268,30 +269,44 @@ export function checkUnsavedChanges(): boolean {
     return vscode.workspace.textDocuments.some((doc) => doc.isDirty)
 }
 
+export const DefaultExtraExcludePatterns = [
+    '**/package-lock.json',
+    '**/yarn.lock',
+    '**/*.zip',
+    '**/*.tar.gz',
+    '**/*.bin',
+    '**/*.png',
+    '**/*.jpg',
+    '**/*.svg',
+    '**/*.pyc',
+    '**/*.pdf',
+    '**/*.ttf',
+    '**/*.ico',
+    '**/license.txt',
+    '**/License.txt',
+    '**/LICENSE.txt',
+    '**/license.md',
+    '**/License.md',
+    '**/LICENSE.md',
+]
+
 export function getExcludePattern(defaultExcludePatterns: boolean = true) {
-    const globAlwaysExcludedDirs = getGlobDirExcludedPatterns().map((pattern) => `**/${pattern}/*`)
-    const extraPatterns = [
-        '**/package-lock.json',
-        '**/yarn.lock',
-        '**/*.zip',
-        '**/*.tar.gz',
-        '**/*.bin',
-        '**/*.png',
-        '**/*.jpg',
-        '**/*.svg',
-        '**/*.pyc',
-        '**/*.pdf',
-        '**/*.ttf',
-        '**/*.ico',
-        '**/license.txt',
-        '**/License.txt',
-        '**/LICENSE.txt',
-        '**/license.md',
-        '**/License.md',
-        '**/LICENSE.md',
-    ]
-    const allPatterns = [...globAlwaysExcludedDirs, ...(defaultExcludePatterns ? extraPatterns : [])]
-    return `{${allPatterns.join(',')}}`
+    const globAlwaysExcludedDirs = getGlobalExcludePatterns()
+    const allPatterns = [...globAlwaysExcludedDirs]
+
+    if (defaultExcludePatterns) {
+        allPatterns.push(...DefaultExtraExcludePatterns)
+    }
+
+    return excludePatternsAsString(allPatterns)
+}
+
+export function getGlobalExcludePatterns() {
+    return getGlobDirExcludedPatterns().map((pattern) => `**/${pattern}/*`)
+}
+
+export function excludePatternsAsString(patterns: string[]): string {
+    return `{${patterns.join(',')}}`
 }
 
 /**
@@ -313,30 +328,33 @@ async function filterOutGitignoredFiles(
     return gitIgnoreFilter.filterFiles(files)
 }
 
+export type CollectFilesResultItem = {
+    workspaceFolder: vscode.WorkspaceFolder
+    relativeFilePath: string
+    fileUri: vscode.Uri
+    fileContent: string
+    zipFilePath: string
+}
+export type CollectFilesFileFilter = (relativePath: string) => boolean // returns true if file should be filtered out
+
 /**
- * collects all files that are marked as source
+ * search and collect source files
  * @param sourcePaths the paths where collection starts
  * @param workspaceFolders the current workspace folders opened
- * @param respectGitIgnore whether to respect gitignore file
- * @param addExtraIgnorePatterns whether to add extra exclude patterns even if not in gitignore
+ * @param options - filtering options
  * @returns all matched files
  */
 export async function collectFiles(
     sourcePaths: string[],
     workspaceFolders: CurrentWsFolders,
-    respectGitIgnore: boolean = true,
-    maxSize = 200 * 1024 * 1024, // 200 MB
-    defaultExcludePatterns: boolean = true
-): Promise<
-    {
-        workspaceFolder: vscode.WorkspaceFolder
-        relativeFilePath: string
-        fileUri: vscode.Uri
-        fileContent: string
-        zipFilePath: string
-    }[]
-> {
-    const storage: Awaited<ReturnType<typeof collectFiles>> = []
+    options?: {
+        maxSizeLimitInBytes?: number // 200 MB default
+        useGitIgnoreFileAsFilter?: boolean // default true
+        extraExcludeFilePatterns?: string[] // default DefaultExtraExcludePatterns
+        extraFileFilterFn?: CollectFilesFileFilter
+    }
+): Promise<CollectFilesResultItem[]> {
+    const storage: Awaited<CollectFilesResultItem[]> = []
 
     const workspaceFoldersMapping = getWorkspaceFoldersByPrefixes(workspaceFolders)
     const workspaceToPrefix = new Map<vscode.WorkspaceFolder, string>(
@@ -360,24 +378,37 @@ export async function collectFiles(
     }
 
     let totalSizeBytes = 0
+
+    const useGitIgnoreFileAsFilter = options?.useGitIgnoreFileAsFilter ?? true
+    const extraExcludeFilePatterns = options?.extraExcludeFilePatterns ?? DefaultExtraExcludePatterns
+    const maxSizeLimitInBytes = options?.maxSizeLimitInBytes ?? maxRepoSizeBytes
+
+    const excludePatterns = [...getGlobalExcludePatterns()]
+    if (extraExcludeFilePatterns.length) {
+        excludePatterns.push(...extraExcludeFilePatterns)
+    }
+    const excludePatternFilter = excludePatternsAsString(excludePatterns)
+
     for (const rootPath of sourcePaths) {
         const allFiles = await vscode.workspace.findFiles(
             new vscode.RelativePattern(rootPath, '**'),
-            getExcludePattern(defaultExcludePatterns)
+            excludePatternFilter
         )
 
-        const files = respectGitIgnore
-            ? await filterOutGitignoredFiles(rootPath, allFiles, defaultExcludePatterns)
-            : allFiles
+        const files = useGitIgnoreFileAsFilter ? await filterOutGitignoredFiles(rootPath, allFiles, false) : allFiles
 
         for (const file of files) {
             const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
             if (!relativePath) {
                 continue
             }
 
+            if (options?.extraFileFilterFn && options.extraFileFilterFn(relativePath.relativePath)) {
+                continue
+            }
+
             const fileStat = await fs.stat(file)
-            if (totalSizeBytes + fileStat.size > maxSize) {
+            if (totalSizeBytes + fileStat.size > maxSizeLimitInBytes) {
                 throw new ToolkitError(
                     'The project you have selected for source code is too large to use as context. Please select a different folder to use',
                     { code: 'ContentLengthError' }