fix(codewhisperer): support code generation for IaC languages without external plugin extensions (#4181)

* CW: Support IaaC languages with file extension instead of editor language for code generation

* Change in test case file

* Added change log

* minor code changes

* Small code fixes

* Addressing comments

* CW:Change log

* Change in text

---------

Co-authored-by: Laxman Reddy <[email protected]>

Author: laileni-aws

.changes/next-release/Bug Fix-ec18555e-d716-40b5-b8d1-54dad851fdc2.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "CodeWhisperer supports code generation for Iac languages without external extensions"
+}

src/codewhisperer/commands/onAcceptance.ts

@@ -14,6 +14,7 @@ import { handleExtraBrackets } from '../util/closingBracketUtil'
 import { RecommendationHandler } from '../service/recommendationHandler'
 import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider'
 import { ReferenceHoverProvider } from '../service/referenceHoverProvider'
+import path from 'path'
 
 /**
  * This function is called when user accepts a intelliSense suggestion or an inline suggestion
@@ -24,7 +25,10 @@ export async function onAcceptance(acceptanceEntry: OnRecommendationAcceptanceEn
      * Format document
      */
     if (acceptanceEntry.editor) {
-        const languageContext = runtimeLanguageContext.getLanguageContext(acceptanceEntry.editor.document.languageId)
+        const languageContext = runtimeLanguageContext.getLanguageContext(
+            acceptanceEntry.editor.document.languageId,
+            path.extname(acceptanceEntry.editor.document.fileName)
+        )
         const start = acceptanceEntry.range.start
         const end = isCloud9() ? acceptanceEntry.editor.selection.active : acceptanceEntry.range.end
 

src/codewhisperer/commands/onInlineAcceptance.ts

@@ -27,6 +27,7 @@ import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider'
 import { ReferenceHoverProvider } from '../service/referenceHoverProvider'
 import { ImportAdderProvider } from '../service/importAdderProvider'
 import { session } from '../util/codeWhispererSession'
+import path from 'path'
 
 export const acceptSuggestion = Commands.declare(
     'aws.codeWhisperer.accept',
@@ -73,7 +74,10 @@ export async function onInlineAcceptance(
 
     if (acceptanceEntry.editor) {
         await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
-        const languageContext = runtimeLanguageContext.getLanguageContext(acceptanceEntry.editor.document.languageId)
+        const languageContext = runtimeLanguageContext.getLanguageContext(
+            acceptanceEntry.editor.document.languageId,
+            path.extname(acceptanceEntry.editor.document.fileName)
+        )
         const start = acceptanceEntry.range.start
         const end = acceptanceEntry.editor.selection.active
 

src/codewhisperer/commands/startSecurityScan.ts

@@ -33,6 +33,7 @@ import { isAwsError } from '../../shared/errors'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { AuthUtil } from '../util/authUtil'
 import { DependencyGraphConstants } from '../util/dependencyGraph/dependencyGraph'
+import path from 'path'
 
 const performance = globalThis.performance ?? require('perf_hooks').performance
 const securityScanOutputChannel = vscode.window.createOutputChannel('CodeWhisperer Security Scan Logs')
@@ -72,7 +73,10 @@ export async function startSecurityScan(
     const codeScanStartTime = performance.now()
     let serviceInvocationStartTime = 0
     const codeScanTelemetryEntry: CodeScanTelemetryEntry = {
-        codewhispererLanguage: runtimeLanguageContext.getLanguageContext(editor.document.languageId).language,
+        codewhispererLanguage: runtimeLanguageContext.getLanguageContext(
+            editor.document.languageId,
+            path.extname(editor.document.fileName)
+        ).language,
         codewhispererCodeScanSrcPayloadBytes: 0,
         codewhispererCodeScanSrcZipFileBytes: 0,
         codewhispererCodeScanLines: 0,

src/codewhisperer/service/completionProvider.ts

@@ -10,6 +10,7 @@ import { Recommendation } from '../client/codewhisperer'
 import { LicenseUtil } from '../util/licenseUtil'
 import { RecommendationHandler } from './recommendationHandler'
 import { session } from '../util/codeWhispererSession'
+import path from 'path'
 /**
  * completion provider for intelliSense popup
  */
@@ -41,7 +42,10 @@ export function getCompletionItem(
     completionItem.preselect = true
     completionItem.sortText = String(recommendationIndex + 1).padStart(10, '0')
     completionItem.range = new vscode.Range(start, position)
-    const languageContext = runtimeLanguageContext.getLanguageContext(document.languageId)
+    const languageContext = runtimeLanguageContext.getLanguageContext(
+        document.languageId,
+        path.extname(document.fileName)
+    )
     let references: typeof recommendationDetail.references
     if (recommendationDetail.references !== undefined && recommendationDetail.references.length > 0) {
         references = recommendationDetail.references

src/codewhisperer/service/inlineCompletionItemProvider.ts

@@ -11,6 +11,7 @@ import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { ReferenceInlineProvider } from './referenceInlineProvider'
 import { ImportAdderProvider } from './importAdderProvider'
 import { application } from '../util/codeWhispererApplication'
+import path from 'path'
 
 export class CWInlineCompletionItemProvider implements vscode.InlineCompletionItemProvider {
     private activeItemIndex: number | undefined
@@ -116,7 +117,8 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
                     session.sessionId,
                     session.triggerType,
                     session.getCompletionType(index),
-                    runtimeLanguageContext.getLanguageContext(document.languageId).language,
+                    runtimeLanguageContext.getLanguageContext(document.languageId, path.extname(document.fileName))
+                        .language,
                     r.references,
                 ],
             },

src/codewhisperer/service/recommendationHandler.ts

@@ -43,6 +43,7 @@ import { CWInlineCompletionItemProvider } from './inlineCompletionItemProvider'
 import { application } from '../util/codeWhispererApplication'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { indent } from '../../shared/utilities/textUtilities'
+import path from 'path'
 
 /**
  * This class is for getRecommendation/listRecommendation API calls and its states
@@ -180,7 +181,10 @@ export class RecommendationHandler {
         let latency = 0
         let nextToken = ''
         let shouldRecordServiceInvocation = true
-        session.language = runtimeLanguageContext.getLanguageContext(editor.document.languageId).language
+        session.language = runtimeLanguageContext.getLanguageContext(
+            editor.document.languageId,
+            path.extname(editor.document.fileName)
+        ).language
         session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName)
 
         if (pagination) {
@@ -399,7 +403,10 @@ export class RecommendationHandler {
                     session.requestIdList,
                     sessionId,
                     page,
-                    editor.document.languageId,
+                    runtimeLanguageContext.getLanguageContext(
+                        editor.document.languageId,
+                        path.extname(editor.document.fileName)
+                    ).language,
                     session.requestContext.supplementalMetadata
                 )
             }
@@ -685,7 +692,10 @@ export class RecommendationHandler {
     private sendPerceivedLatencyTelemetry() {
         if (vscode.window.activeTextEditor) {
             const languageContext = runtimeLanguageContext.getLanguageContext(
-                vscode.window.activeTextEditor.document.languageId
+                vscode.window.activeTextEditor.document.languageId,
+                vscode.window.activeTextEditor.document.fileName.substring(
+                    vscode.window.activeTextEditor.document.fileName.lastIndexOf('.') + 1
+                )
             )
             telemetry.codewhisperer_perceivedLatency.emit({
                 codewhispererRequestId: this.requestId,

src/codewhisperer/util/editorContext.ts

@@ -181,12 +181,14 @@ export async function buildGenerateRecommendationRequest(editor: vscode.TextEdit
 export function validateRequest(
     req: codewhispererClient.ListRecommendationsRequest | codewhispererClient.GenerateRecommendationsRequest
 ): boolean {
-    const isLanguageNameValid = !(
-        req.fileContext.programmingLanguage.languageName === undefined ||
-        req.fileContext.programmingLanguage.languageName.length < 1 ||
-        req.fileContext.programmingLanguage.languageName.length > 128 ||
-        !runtimeLanguageContext.isLanguageSupported(req.fileContext.programmingLanguage.languageName)
-    )
+    const isLanguageNameValid =
+        req.fileContext.programmingLanguage.languageName !== undefined &&
+        req.fileContext.programmingLanguage.languageName.length >= 1 &&
+        req.fileContext.programmingLanguage.languageName.length <= 128 &&
+        (runtimeLanguageContext.isLanguageSupported(req.fileContext.programmingLanguage.languageName) ||
+            runtimeLanguageContext.isFileFormatSupported(
+                req.fileContext.filename.substring(req.fileContext.filename.lastIndexOf('.') + 1)
+            ))
     const isFileNameValid = !(req.fileContext.filename === undefined || req.fileContext.filename.length < 1)
     const isFileContextValid = !(
         req.fileContext.leftFileContent.length > CodeWhispererConstants.charactersLimit ||

src/codewhisperer/util/runtimeLanguageContext.ts

@@ -158,10 +158,25 @@ export class RuntimeLanguageContext {
     }
 
     /**
-     * @param languageId : vscode language id or codewhisperer language name
+     * @param languageId : vscode language id or codewhisperer language name, fileExtension: extension of the selected file
      * @returns An object with a field language: CodewhispererLanguage, if no corresponding CodewhispererLanguage ID, plaintext is returned
      */
-    public getLanguageContext(languageId?: string): { language: CodewhispererLanguage } {
+    public getLanguageContext(languageId?: string, fileExtension?: string): { language: CodewhispererLanguage } {
+        const extensionToLanguageMap: Record<string, CodewhispererLanguage> = {
+            tf: 'tf',
+            hcl: 'tf',
+            json: 'json',
+            yaml: 'yaml',
+            yml: 'yaml',
+            // Add more mappings if needed
+        }
+
+        if (languageId === 'plaintext' && fileExtension !== undefined) {
+            const languages = extensionToLanguageMap[fileExtension]
+            if (languages) {
+                return { language: languages }
+            }
+        }
         return { language: this.normalizeLanguage(languageId) ?? 'plaintext' }
     }
 
@@ -177,7 +192,10 @@ export class RuntimeLanguageContext {
         const fileContext = request.fileContext
         const runtimeLanguage: codewhispererClient.ProgrammingLanguage = {
             languageName: this.toRuntimeLanguage(
-                request.fileContext.programmingLanguage.languageName as CodewhispererLanguage
+                runtimeLanguageContext.getLanguageContext(
+                    request.fileContext.programmingLanguage.languageName,
+                    request.fileContext.filename.substring(request.fileContext.filename.lastIndexOf('.') + 1)
+                ).language
             ),
         }
 
@@ -193,12 +211,21 @@ export class RuntimeLanguageContext {
     /**
      *
      * @param languageId: either vscodeLanguageId or CodewhispererLanguage
-     * @returns ture if the language is supported by CodeWhisperer otherwise false
+     * @returns true if the language is supported by CodeWhisperer otherwise false
      */
     public isLanguageSupported(languageId: string): boolean {
         const lang = this.normalizeLanguage(languageId)
         return lang !== undefined && this.normalizeLanguage(languageId) !== 'plaintext'
     }
+    /**
+     *
+     * @param fileFormat : vscode editor filecontext filename extension
+     * @returns  true if the fileformat is supported by CodeWhisperer otherwise false
+     */
+    public isFileFormatSupported(fileFormat: string): boolean {
+        const fileformat = this.supportedLanguageExtensionMap.get(fileFormat)
+        return fileformat !== undefined && this.supportedLanguageExtensionMap.get(fileformat) !== 'txt'
+    }
 }
 
 export const runtimeLanguageContext = new RuntimeLanguageContext()

src/codewhisperer/util/telemetryHelper.ts

@@ -99,10 +99,9 @@ export class TelemetryHelper {
         requestIdList: string[],
         sessionId: string,
         paginationIndex: number,
-        languageId: string,
+        language: CodewhispererLanguage,
         supplementalContextMetadata?: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
     ) {
-        const languageContext = runtimeLanguageContext.getLanguageContext(languageId)
         telemetry.codewhisperer_userDecision.emit({
             codewhispererRequestId: requestIdList[0],
             codewhispererSessionId: sessionId ? sessionId : undefined,
@@ -113,7 +112,7 @@ export class TelemetryHelper {
             codewhispererSuggestionReferences: undefined,
             codewhispererSuggestionReferenceCount: 0,
             codewhispererCompletionType: 'Line',
-            codewhispererLanguage: languageContext.language,
+            codewhispererLanguage: language,
             codewhispererGettingStartedTask: session.taskType,
             credentialStartUrl: AuthUtil.instance.startUrl,
             codewhispererUserGroup: CodeWhispererUserGroupSettings.getUserGroup().toString(),