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/amazonq/test/unit/amazonqFeatureDev/util/files.test.ts

@@ -6,6 +6,7 @@ import * as vscode from 'vscode'
 import assert from 'assert'
 import {
     prepareRepoData,
+    PrepareRepoDataOptions,
     TelemetryHelper,
     ContentLengthError,
     maxRepoSizeBytes,
@@ -44,19 +45,24 @@ const testDevfilePrepareRepo = async (devfileEnabled: boolean) => {
         .stub(CodeWhispererSettings.instance, 'getAutoBuildSetting')
         .returns(devfileEnabled ? { [workspace.uri.fsPath]: true } : {})
 
-    await testPrepareRepoData(workspace, expectedFiles)
+    await testPrepareRepoData(workspace, expectedFiles, { telemetry: new TelemetryHelper() })
 }
 
 const testPrepareRepoData = async (
     workspace: vscode.WorkspaceFolder,
     expectedFiles: string[],
+    prepareRepoDataOptions: PrepareRepoDataOptions,
     expectedTelemetryMetrics?: Array<{ metricName: MetricName; value: any }>
 ) => {
     expectedFiles.sort((a, b) => a.localeCompare(b))
-    const telemetry = new TelemetryHelper()
-    const result = await prepareRepoData([workspace.uri.fsPath], [workspace], telemetry, {
-        record: () => {},
-    } as unknown as Span<AmazonqCreateUpload>)
+    const result = await prepareRepoData(
+        [workspace.uri.fsPath],
+        [workspace],
+        {
+            record: () => {},
+        } as unknown as Span<AmazonqCreateUpload>,
+        prepareRepoDataOptions
+    )
 
     assert.strictEqual(Buffer.isBuffer(result.zipFileBuffer), true)
     // checksum is not the same across different test executions because some unique random folder names are generated
@@ -77,6 +83,8 @@ const testPrepareRepoData = async (
 
 describe('file utils', () => {
     describe('prepareRepoData', function () {
+        const defaultPrepareRepoDataOptions: PrepareRepoDataOptions = { telemetry: new TelemetryHelper() }
+
         afterEach(() => {
             sinon.restore()
         })
@@ -85,21 +93,33 @@ describe('file utils', () => {
             const folder = await TestFolder.create()
             await folder.write('file1.md', 'test content')
             await folder.write('file2.md', 'test content')
+            await folder.write('docs/infra.svg', 'test content')
+            const workspace = getWorkspaceFolder(folder.path)
+
+            await testPrepareRepoData(workspace, ['file1.md', 'file2.md'], defaultPrepareRepoDataOptions)
+        })
+
+        it('infrastructure diagram is included', async function () {
+            const folder = await TestFolder.create()
+            await folder.write('file1.md', 'test content')
+            await folder.write('file2.svg', 'test content')
+            await folder.write('docs/infra.svg', 'test content')
             const workspace = getWorkspaceFolder(folder.path)
 
-            await testPrepareRepoData(workspace, ['file1.md', 'file2.md'])
+            await testPrepareRepoData(workspace, ['file1.md', 'docs/infra.svg'], {
+                telemetry: new TelemetryHelper(),
+                isIncludeInfraDiagram: true,
+            })
         })
 
         it('prepareRepoData ignores denied file extensions', async function () {
             const folder = await TestFolder.create()
             await folder.write('file.mp4', 'test content')
             const workspace = getWorkspaceFolder(folder.path)
 
-            await testPrepareRepoData(
-                workspace,
-                [],
-                [{ metricName: 'amazonq_bundleExtensionIgnored', value: { filenameExt: 'mp4', count: 1 } }]
-            )
+            await testPrepareRepoData(workspace, [], defaultPrepareRepoDataOptions, [
+                { metricName: 'amazonq_bundleExtensionIgnored', value: { filenameExt: 'mp4', count: 1 } },
+            ])
         })
 
         it('should ignore devfile.yaml when setting is disabled', async function () {
@@ -115,14 +135,18 @@ describe('file utils', () => {
             const folder = await TestFolder.create()
             await folder.write('file.md', 'test content')
             const workspace = getWorkspaceFolder(folder.path)
-            const telemetry = new TelemetryHelper()
 
             sinon.stub(fs, 'stat').resolves({ size: 2 * maxRepoSizeBytes } as vscode.FileStat)
             await assert.rejects(
                 () =>
-                    prepareRepoData([workspace.uri.fsPath], [workspace], telemetry, {
-                        record: () => {},
-                    } as unknown as Span<AmazonqCreateUpload>),
+                    prepareRepoData(
+                        [workspace.uri.fsPath],
+                        [workspace],
+                        {
+                            record: () => {},
+                        } as unknown as Span<AmazonqCreateUpload>,
+                        defaultPrepareRepoDataOptions
+                    ),
                 ContentLengthError
             )
         })

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'
@@ -25,7 +25,7 @@ import {
     SessionStateInteraction,
     SessionStatePhase,
 } from '../commons/types'
-import { prepareRepoData, getDeletedFileInfos, registerNewFiles } from '../util/files'
+import { prepareRepoData, getDeletedFileInfos, registerNewFiles, PrepareRepoDataOptions } from '../util/files'
 import { uploadCode } from '../util/upload'
 
 export const EmptyCodeGenID = 'EMPTY_CURRENT_CODE_GENERATION_ID'
@@ -227,11 +227,11 @@ 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,
-                span
+                span,
+                { telemetry: action.telemetry }
             )
             const uploadId = randomUUID()
             const { uploadUrl, kmsKeyArn } = await this.config.proxyClient.createUploadUrl(
@@ -251,6 +251,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,
+        span: Span<AmazonqCreateUpload>,
+        options: PrepareRepoDataOptions
+    ) {
+        return await prepareRepoData(workspaceRoots, workspaceFolders, span, options)
+    }
 }
 
 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'
@@ -24,27 +29,71 @@ import { isPresent } from '../../shared/utilities/collectionUtils'
 import { AuthUtil } from '../../codewhisperer/util/authUtil'
 import { TelemetryHelper } from '../util/telemetryHelper'
 
+export const SvgFileExtension = '.svg'
+
 export async function checkForDevFile(root: string) {
     const devFilePath = root + '/devfile.yaml'
     const hasDevFile = await fs.existsFile(devFilePath)
     return hasDevFile
 }
 
+function isInfraDiagramFile(relativePath: string) {
+    return (
+        relativePath.toLowerCase().endsWith(path.join('docs', 'infra.dot')) ||
+        relativePath.toLowerCase().endsWith(path.join('docs', 'infra.svg'))
+    )
+}
+
+export type PrepareRepoDataOptions = {
+    telemetry?: TelemetryHelper
+    zip?: ZipStream
+    isIncludeInfraDiagram?: boolean
+}
+
 /**
  * given the root path of the repo it zips its files in memory and generates a checksum for it.
  */
 export async function prepareRepoData(
     repoRootPaths: string[],
     workspaceFolders: CurrentWsFolders,
-    telemetry: TelemetryHelper,
     span: Span<AmazonqCreateUpload>,
-    zip: ZipStream = new ZipStream()
+    options?: PrepareRepoDataOptions
 ) {
     try {
+        const telemetry = options?.telemetry
+        const isIncludeInfraDiagram = options?.isIncludeInfraDiagram ?? false
+        const zip = options?.zip ?? new ZipStream()
+
         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(SvgFileExtension))
+                )
+                // ensure only infra diagram is included from all svg files
+                extraFileFilterFn = (relativePath: string) => {
+                    if (!relativePath.toLowerCase().endsWith(SvgFileExtension)) {
+                        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 +111,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 +132,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
@@ -111,7 +167,10 @@ export async function prepareRepoData(
             }
         }
 
-        telemetry.setRepositorySize(totalBytes)
+        if (telemetry) {
+            telemetry.setRepositorySize(totalBytes)
+        }
+
         span.record({ amazonqRepositorySize: totalBytes })
         const zipResult = await zip.finalize()
 

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

@@ -35,13 +35,13 @@ 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,
     isMultiRootWorkspace,
 } from '../../../shared/utilities/workspaceUtils'
-import { getPathsFromZipFilePath } from '../../../amazonq/util/files'
+import { getPathsFromZipFilePath, SvgFileExtension } from '../../../amazonq/util/files'
 import { FollowUpTypes } from '../../../amazonq/commons/types'
 import { DocGenerationTask } from '../docGenerationTask'
 import { DevPhase } from '../../types'
@@ -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(SvgFileExtension)) {
+                    const rightPathUri = createAmazonQUri(rightPath, tabId, this.scheme)
+                    const infraDiagramContent = await vscode.workspace.openTextDocument(rightPathUri)
+                    await vscode.window.showTextDocument(infraDiagramContent)
+                } 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,8 @@ import {
 import { DocMessenger } from '../messenger'
 import { BaseCodeGenState, BasePrepareCodeGenState, CreateNextStateParams } from '../../amazonq/session/sessionState'
 import { Intent } from '../../amazonq/commons/types'
+import { AmazonqCreateUpload, Span } from '../../shared/telemetry/telemetry'
+import { prepareRepoData, PrepareRepoDataOptions } from '../../amazonq/util/files'
 
 export class DocCodeGenState extends BaseCodeGenState {
     protected handleProgress(messenger: DocMessenger, action: SessionStateAction, detail?: string): void {
@@ -132,4 +134,16 @@ export class DocPrepareCodeGenState extends BasePrepareCodeGenState {
     protected override createNextState(config: SessionStateConfig): SessionState {
         return super.createNextState(config, DocCodeGenState)
     }
+
+    protected override async prepareProjectZip(
+        workspaceRoots: string[],
+        workspaceFolders: CurrentWsFolders,
+        span: Span<AmazonqCreateUpload>,
+        options: PrepareRepoDataOptions
+    ) {
+        return await prepareRepoData(workspaceRoots, workspaceFolders, span, {
+            ...options,
+            isIncludeInfraDiagram: true,
+        })
+    }
 }

packages/core/src/amazonqFeatureDev/index.ts

@@ -11,5 +11,5 @@ export { Session } from './session/session'
 export { FeatureDevClient } from './client/featureDev'
 export { FeatureDevChatSessionStorage } from './storages/chatSession'
 export { TelemetryHelper } from '../amazonq/util/telemetryHelper'
-export { prepareRepoData } from '../amazonq/util/files'
+export { prepareRepoData, PrepareRepoDataOptions } from '../amazonq/util/files'
 export { ChatControllerEventEmitters, FeatureDevController } from './controllers/chat/controller'

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)